import { app, safeStorage } from 'electron'
import { promises as fs } from 'node:fs'
import { join, isAbsolute } from 'node:path'
import { createHash, randomUUID } from 'node:crypto'
import { Client } from '@modelcontextprotocol/sdk/client/index.js'
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'
import type { McpServerInput, McpServerView, McpToolView, McpStatus } from '@shared/capabilities'
import type { ToolSpec } from '@shared/types'

interface SecretValues { env: Record<string, string>; headers: Record<string, string> }
interface StoredServer extends Omit<McpServerView, 'status' | 'error' | 'trusted'> {
  secrets?: string
  trustedHash?: string
}
interface Registry { version: 1; servers: StoredServer[] }
interface Runtime {
  generation: string
  configHash: string
  status: McpStatus
  error?: string
  client: Client
  transport: Transport
  controller: AbortController
  secretValues: string[]
  flight?: Promise<McpServerView>
}
let registry: Registry | null = null
let loading: Promise<Registry> | null = null
let writes: Promise<unknown> = Promise.resolve()
const runtimes = new Map<string, Runtime>()
const MAX_SERVERS = 100, MAX_TOOLS = 10000, CONNECT_TIMEOUT = 20000, CALL_TIMEOUT = 120000
const filePath = (): string => join(app.getPath('userData'), 'mcp-servers.json')
const clone = <T>(value: T): T => structuredClone(value)
const abortError = (): Error => Object.assign(new Error('MCP 调用已取消'), { name: 'AbortError' })
function object(value: unknown): value is Record<string, unknown> { return !!value && typeof value === 'object' && !Array.isArray(value) }
function sorted(value: unknown): unknown {
  if (Array.isArray(value)) return value.map(sorted)
  return object(value) ? Object.fromEntries(Object.keys(value).sort().map(key => [key, sorted(value[key])])) : value
}
function readSecrets(server: StoredServer): SecretValues {
  if (!server.secrets) return { env: {}, headers: {} }
  try {
    const values = JSON.parse(safeStorage.decryptString(Buffer.from(server.secrets, 'base64'))) as SecretValues
    if (!object(values.env) || !object(values.headers)) throw Error()
    return values
  } catch { throw new Error('无法解密 MCP 凭据，请重新配置') }
}
function encryptSecrets(values: SecretValues): string | undefined {
  if (!Object.keys(values.env).length && !Object.keys(values.headers).length) return undefined
  try {
    if (!safeStorage.isEncryptionAvailable()) throw Error()
    return safeStorage.encryptString(JSON.stringify(values)).toString('base64')
  } catch { throw new Error('系统加密服务不可用，无法安全保存 MCP 凭据') }
}
function mergeSecretMap(input: unknown, previous: Record<string, string>, headers: boolean): Record<string, string> {
  if (input === undefined) return { ...previous }
  if (!object(input) || Object.keys(input).length > 128) throw new Error('MCP 环境变量或 Headers 格式无效')
  const result: Record<string, string> = Object.create(null)
  for (const [rawKey, value] of Object.entries(input)) {
    const key = headers ? rawKey.toLowerCase() : rawKey
    if (!key || key.length > 256 || !/^[A-Za-z_][A-Za-z0-9_.-]*$/.test(key) || typeof value !== 'string' || value.length > 65536) throw new Error('MCP 环境变量或 Headers 格式无效')
    if (headers && /[\r\n]/.test(value)) throw new Error('MCP Headers 格式无效')
    // Empty retained keys keep credentials; omit a key from an explicitly supplied map to remove it.
    if (value) result[key] = value
    else if (previous[key] !== undefined) result[key] = previous[key]
  }
  return result
}
function normalize(input: McpServerInput, previous?: StoredServer): StoredServer {
  if (!input || typeof input.name !== 'string' || !input.name.trim() || input.name.length > 128) throw new Error('请输入有效的 MCP 名称')
  if (input.id && (!previous || previous.id !== input.id)) throw new Error('MCP 服务不存在')
  if (input.transport !== 'stdio' && input.transport !== 'http') throw new Error('不支持的 MCP 连接类型')
  const before = previous ? readSecrets(previous) : { env: {}, headers: {} }
  const secrets = { env: mergeSecretMap(input.env, before.env, false), headers: mergeSecretMap(input.headers, before.headers, true) }
  const config: Omit<McpServerInput, 'id' | 'env' | 'headers' | 'enabled'> = { name: input.name.trim(), transport: input.transport }
  if (input.transport === 'stdio') {
    if (typeof input.command !== 'string' || !input.command.trim() || input.command.length > 4096 || /[\r\n\0]/.test(input.command)) throw new Error('请输入有效的 MCP 启动命令')
    if (input.args !== undefined && (!Array.isArray(input.args) || input.args.length > 256 || input.args.some(arg => typeof arg !== 'string' || arg.length > 65536 || arg.includes('\0')))) throw new Error('MCP 参数格式无效')
    if (input.cwd && (!isAbsolute(input.cwd) || input.cwd.includes('\0'))) throw new Error('MCP 工作目录必须为绝对路径')
    config.command = input.command.trim(); config.args = input.args?.slice() || []
    if (input.cwd) config.cwd = input.cwd
  } else {
    let url: URL
    try { url = new URL(input.url || '') } catch { throw new Error('请输入有效的 MCP HTTP 地址') }
    if (!['http:', 'https:'].includes(url.protocol) || url.username || url.password || [...url.searchParams.keys()].some(key => /token|key|secret|password|authorization/i.test(key))) throw new Error('MCP 地址必须为 HTTP(S)，凭据请使用 Headers 配置')
    config.url = url.href
  }
  const configHash = createHash('sha256').update(JSON.stringify(sorted({ ...config, ...secrets }))).digest('hex')
  return { ...config, id: previous?.id || randomUUID(), enabled: input.enabled ?? previous?.enabled ?? false,
    envKeys: Object.keys(secrets.env).sort(), headerKeys: Object.keys(secrets.headers).sort(), configHash,
    secrets: encryptSecrets(secrets), trustedHash: previous?.configHash === configHash ? previous.trustedHash : undefined,
    tools: previous?.configHash === configHash ? clone(previous.tools) : [] }
}
async function load(): Promise<Registry> {
  if (registry) return registry
  if (loading) return loading
  loading = (async () => {
    let content: string
    try { content = await fs.readFile(filePath(), 'utf8') }
    catch (error) { if ((error as NodeJS.ErrnoException).code === 'ENOENT') return registry = { version: 1, servers: [] }; throw error }
    let parsed: Registry
    try { parsed = JSON.parse(content) as Registry } catch { throw new Error('MCP 配置文件损坏，请保留配置文件后检查内容') }
    if (parsed.version !== 1 || !Array.isArray(parsed.servers) || parsed.servers.length > MAX_SERVERS) throw new Error('MCP 配置版本或内容无效')
    const ids = new Set<string>()
    const normalized: StoredServer[] = []
    for (const server of parsed.servers) {
      if (!server || typeof server.id !== 'string' || ids.has(server.id) || typeof server.name !== 'string' || typeof server.configHash !== 'string'
        || !['stdio', 'http'].includes(server.transport) || !Array.isArray(server.tools) || !Array.isArray(server.envKeys) || !Array.isArray(server.headerKeys)) throw new Error('MCP 配置内容无效')
      ids.add(server.id)
      // A hand-edited file cannot retain approval for a different command, URL or secret.
      // Reconstructing the record also excludes unknown credential-bearing fields from IPC views.
      normalized.push(normalize({ ...server, ...readSecrets(server) }, server))
    }
    return registry = { version: 1, servers: normalized }
  })().finally(() => { loading = null })
  return loading
}
async function persist(next: Registry): Promise<void> {
  const file = filePath(), temporary = `${file}.${randomUUID()}.tmp`
  await fs.mkdir(join(file, '..'), { recursive: true })
  try { await fs.writeFile(temporary, JSON.stringify(next, null, 2), { encoding: 'utf8', mode: 0o600 }); await fs.rename(temporary, file); registry = next }
  finally { await fs.unlink(temporary).catch(() => undefined) }
}
function mutate<T>(operation: (current: Registry) => Promise<T>): Promise<T> {
  const next = writes.then(async () => operation(await load()))
  writes = next.catch(() => undefined)
  return next
}
function requireServer(current: Registry, id: string): StoredServer {
  const server = current.servers.find(server => server.id === id)
  if (!server) throw new Error('MCP 服务不存在')
  return server
}
function view(server: StoredServer): McpServerView {
  const { secrets: _secret, trustedHash, ...publicFields } = server
  const runtime = runtimes.get(server.id)
  return { ...clone(publicFields), trusted: trustedHash === server.configHash,
    status: runtime?.status || 'disconnected', ...(runtime?.error ? { error: runtime.error } : {}) }
}
function publicError(error: unknown, operation: 'connect' | 'call'): string {
  const message = error instanceof Error ? error.message : ''
  if (/401|403|unauthoriz|forbidden/i.test(message)) return 'MCP 鉴权失败，请检查凭据与权限'
  if (/timeout|timed out/i.test(message)) return 'MCP 请求超时，请稍后重试'
  return operation === 'connect' ? '无法连接 MCP 服务，请检查命令、地址和运行环境' : 'MCP 工具调用失败'
}
async function closeRuntime(id: string): Promise<void> {
  const runtime = runtimes.get(id)
  if (!runtime) return
  runtimes.delete(id)
  runtime.controller.abort()
  // Close aborts all pending requests before best-effort remote session cleanup.
  if (runtime.transport instanceof StreamableHTTPClientTransport) {
    await Promise.race([runtime.transport.terminateSession().catch(() => undefined), new Promise(resolve => setTimeout(resolve, 500))])
  }
  await runtime.client.close().catch(() => undefined)
}
export async function listMcpServers(): Promise<McpServerView[]> { await writes; return (await load()).servers.map(view) }
export async function saveMcpServer(input: McpServerInput): Promise<McpServerView> {
  return mutate(async current => {
    const previous = input.id ? requireServer(current, input.id) : undefined, next = normalize(input, previous)
    if (!previous && current.servers.length >= MAX_SERVERS) throw new Error('MCP 服务数量已达上限')
    if (current.servers.some(server => server.id !== next.id && server.name === next.name)) throw new Error('同名 MCP 已存在')
    await persist({ version: 1, servers: previous ? current.servers.map(server => server.id === next.id ? next : server) : [...current.servers, next] })
    if (previous && (previous.configHash !== next.configHash || !next.enabled)) await closeRuntime(next.id)
    return view(next)
  })
}
export async function trustMcpServer(id: string, configHash: string): Promise<void> {
  await mutate(async current => {
    const server = requireServer(current, id)
    if (server.configHash !== configHash) throw new Error('MCP 配置已变化，请重新查看并批准')
    await persist({ ...current, servers: current.servers.map(item => item.id === id ? { ...item, trustedHash: configHash } : item) })
  })
}
export async function connectMcpServer(id: string): Promise<McpServerView> {
  await writes
  const server = requireServer(await load(), id)
  if (!server.enabled) throw new Error('MCP 服务已禁用')
  if (server.trustedHash !== server.configHash) throw new Error('请先批准当前 MCP 配置')
  const existing = runtimes.get(id)
  if (existing?.status === 'connected' && existing.configHash === server.configHash) return view(server)
  if (existing?.flight && existing.configHash === server.configHash) return existing.flight
  if (existing) await closeRuntime(id)
  await writes
  const current = requireServer(await load(), id)
  if (current.configHash !== server.configHash || !current.enabled || current.trustedHash !== current.configHash) throw new Error('MCP 配置已变化，请重试')
  const values = readSecrets(server), client = new Client({ name: 'YCode', version: '0.1.0' })
  const transport: Transport = server.transport === 'stdio'
    ? new StdioClientTransport({ command: server.command!, args: server.args, cwd: server.cwd, env: values.env, stderr: 'ignore' })
    : new StreamableHTTPClientTransport(new URL(server.url!), { requestInit: { headers: values.headers, redirect: 'error' },
      // Keep credentials on the configured origin. GET streams also use this guarded fetch.
      fetch: async (input, init) => {
        const target = new URL(input instanceof Request ? input.url : String(input)), configured = new URL(server.url!)
        if (target.origin !== configured.origin) throw new Error('MCP origin changed')
        return fetch(input, { ...init, redirect: 'error' })
      } })
  const runtime: Runtime = { generation: randomUUID(), configHash: server.configHash, status: 'connecting', client, transport,
    controller: new AbortController(), secretValues: [...Object.values(values.env), ...Object.values(values.headers),
      ...Object.values(values.headers).filter(value => /^Bearer\s+/i.test(value)).map(value => value.replace(/^Bearer\s+/i, ''))].filter(Boolean) }
  runtimes.set(id, runtime)
  client.onclose = () => { if (runtimes.get(id) === runtime && runtime.status === 'connected') { runtime.status = 'error'; runtime.error = 'MCP 服务连接已断开' } }
  client.onerror = () => {} // Do not forward upstream payloads/credentials into logs or the renderer.
  runtime.flight = (async () => {
    try {
      await client.connect(transport, { signal: runtime.controller.signal, timeout: CONNECT_TIMEOUT })
      const tools: McpToolView[] = [], names = new Set<string>(), cursors = new Set<string>()
      let cursor: string | undefined
      do {
        const page = await client.listTools(cursor ? { cursor } : undefined, { signal: runtime.controller.signal, timeout: CONNECT_TIMEOUT })
        for (const tool of page.tools) {
          if (names.has(tool.name)) continue
          names.add(tool.name)
          if (tools.length >= MAX_TOOLS) throw Error('MCP tool budget exceeded')
          if (runtime.secretValues.some(secret => tool.name.includes(secret))) throw Error('MCP tool name contains a credential')
          tools.push({ name: tool.name, description: mask(tool.description || '', runtime.secretValues) as string, inputSchema: mask(tool.inputSchema, runtime.secretValues) as Record<string, unknown>,
            enabled: server.tools.find(previous => previous.name === tool.name)?.enabled !== false })
        }
        cursor = page.nextCursor
        if (cursor && cursors.has(cursor)) throw Error('MCP repeated pagination cursor')
        if (cursor) cursors.add(cursor)
      } while (cursor)
      await mutate(async current => {
        const latest = requireServer(current, id)
        if (runtimes.get(id) !== runtime || latest.configHash !== runtime.configHash || !latest.enabled || latest.trustedHash !== latest.configHash) throw abortError()
        // A concurrent tool toggle wins over an older list result.
        const merged = tools.map(tool => ({ ...tool, enabled: latest.tools.find(previous => previous.name === tool.name)?.enabled ?? tool.enabled }))
        await persist({ ...current, servers: current.servers.map(item => item.id === id ? { ...item, tools: merged } : item) })
        runtime.status = 'connected'; runtime.error = undefined
      })
    } catch (error) {
      await client.close().catch(() => undefined)
      if (runtimes.get(id) === runtime) { runtime.status = runtime.controller.signal.aborted ? 'disconnected' : 'error'; runtime.error = runtime.controller.signal.aborted ? undefined : publicError(error, 'connect') }
    } finally { runtime.flight = undefined }
    const latest = (await load()).servers.find(server => server.id === id)
    if (!latest) throw new Error('MCP 服务已移除')
    return view(latest)
  })()
  return runtime.flight
}
export async function disconnectMcpServer(id: string): Promise<void> { await writes; requireServer(await load(), id); await closeRuntime(id) }
export async function setMcpEnabled(id: string, enabled: boolean): Promise<void> {
  await mutate(async current => {
    const server = requireServer(current, id)
    if (enabled && server.trustedHash !== server.configHash) throw new Error('请先批准当前 MCP 配置')
    await persist({ ...current, servers: current.servers.map(item => item.id === id ? { ...item, enabled } : item) })
    if (!enabled) await closeRuntime(id)
  })
  if (enabled) await connectMcpServer(id)
}
export async function setMcpToolEnabled(id: string, toolName: string, enabled: boolean): Promise<void> {
  await mutate(async current => {
    const server = requireServer(current, id)
    if (!server.tools.some(tool => tool.name === toolName)) throw new Error('MCP 工具不存在')
    await persist({ ...current, servers: current.servers.map(item => item.id === id ? { ...item, tools: item.tools.map(tool => tool.name === toolName ? { ...tool, enabled } : tool) } : item) })
  })
}
export async function removeMcpServer(id: string): Promise<void> {
  await mutate(async current => { requireServer(current, id); await persist({ ...current, servers: current.servers.filter(item => item.id !== id) }); await closeRuntime(id) })
}
export async function disposeMcpServers(): Promise<void> { await Promise.all([...runtimes.keys()].map(closeRuntime)) }
/** Startup reconnect is background work; ordinary sends never await unavailable MCP servers. */
export async function initializeMcpServers(): Promise<void> {
  await writes
  for (const server of (await load()).servers) if (server.enabled && server.trustedHash === server.configHash) void connectMcpServer(server.id).catch(() => undefined)
}
function mask(value: unknown, secrets: string[]): unknown {
  if (typeof value === 'string') { for (const secret of secrets) if (secret) value = (value as string).split(secret).join('[已隐藏]'); return value }
  if (Array.isArray(value)) return value.map(item => mask(item, secrets))
  return object(value) ? Object.fromEntries(Object.entries(value).map(([key, item]) => [key, mask(item, secrets)])) : value
}
/** Namespace and connection generation are captured once; changed configurations fail closed. */
export async function snapshotMcpTools(): Promise<{ specs: ToolSpec[]; describe(name: string): { serverName: string; toolName: string } | undefined; execute(name: string, args: Record<string, unknown>, signal: AbortSignal): Promise<{ ok: boolean; content: string; data?: unknown }> }> {
  await writes
  await load()
  const specs: ToolSpec[] = [], mappings = new Map<string, { id: string; serverName: string; tool: string; runtime: Runtime }>()
  for (const server of (await load()).servers) {
    const runtime = runtimes.get(server.id)
    if (!server.enabled || server.trustedHash !== server.configHash || runtime?.status !== 'connected') continue
    for (const tool of server.tools.filter(tool => tool.enabled)) {
      const suffix = createHash('sha256').update(server.id + '\0' + tool.name).digest('hex').slice(0, 16)
      const prefix = tool.name.replace(/[^A-Za-z0-9_-]/g, '_').slice(0, 24) || 'tool', name = `mcp_${prefix}_${suffix}`
      specs.push({ type: 'function', function: { name, description: `[${server.name}] ${tool.description}`.slice(0, 8000), parameters: clone(tool.inputSchema) } })
      mappings.set(name, { id: server.id, serverName: server.name, tool: tool.name, runtime })
    }
  }
  return { specs, describe: name => { const mapping = mappings.get(name); return mapping ? { serverName: mapping.serverName, toolName: mapping.tool } : undefined }, execute: async (name, args, signal) => {
    if (signal.aborted) throw abortError()
    const mapping = mappings.get(name)
    if (!mapping) return { ok: false, content: 'MCP 工具不在本回合快照中' }
    await writes
    const server = (await load()).servers.find(server => server.id === mapping.id), runtime = mapping.runtime
    if (!server || !server.enabled || server.trustedHash !== server.configHash || server.configHash !== runtime.configHash
      || runtimes.get(mapping.id) !== runtime || runtime.status !== 'connected' || !server.tools.some(tool => tool.name === mapping.tool && tool.enabled)) return { ok: false, content: 'MCP 配置或连接已变化，请重新发起请求' }
    try {
      const result = await runtime.client.callTool({ name: mapping.tool, arguments: args }, undefined,
        { signal: AbortSignal.any([signal, runtime.controller.signal]), timeout: CALL_TIMEOUT })
      const data = mask(result, runtime.secretValues) as Record<string, unknown>
      const content = Array.isArray(data.content) ? data.content.map(part => object(part) && part.type === 'text' && typeof part.text === 'string' ? part.text : JSON.stringify(part)).join('\n') : JSON.stringify(data)
      return { ok: !result.isError, content, data }
    } catch (error) {
      if (signal.aborted) throw abortError()
      if (runtime.controller.signal.aborted) return { ok: false, content: 'MCP 配置或连接已变化，请重新发起请求' }
      return { ok: false, content: publicError(error, 'call') }
    }
  } }
}