import assert from 'node:assert/strict'
import { testTerminalHostPlugin } from './test-terminal-host-plugin'
import { mkdir, mkdtemp, readFile, rm } from 'node:fs/promises'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { pathToFileURL } from 'node:url'
import { build } from 'esbuild'
import type { ChatMessage } from '../src/shared/types'

interface PersistenceTestState {
  userData: string
  handlers: Map<string, (...args: unknown[]) => unknown>
  promptCalls: number
  abortCalls: number
  cancelOutputEmitted: boolean
  useOmp: boolean
  fallbackHistory: ChatMessage[] | null
  fallbackWireHistory?: Array<{ role: string; content: string }>
}

const TEST_STATE_KEY = '__YCODE_CHAT_PERSISTENCE_TEST__'

function wait(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms))
}

async function waitForMessages(
  invoke: (channel: string, ...args: unknown[]) => Promise<unknown>,
  sessionId: string,
  done: (messages: ChatMessage[]) => boolean
): Promise<ChatMessage[]> {
  for (let attempt = 0; attempt < 100; attempt++) {
    const messages = (await invoke('session:messages', sessionId)) as ChatMessage[]
    if (done(messages)) return messages
    await wait(10)
  }
  const last = (await invoke('session:messages', sessionId)) as ChatMessage[]
  assert.fail(`Timed out waiting for persisted messages: ${JSON.stringify(last)}`)
}

async function readDiskMessages(userData: string, sessionId: string): Promise<ChatMessage[]> {
  const stored = JSON.parse(await readFile(join(userData, 'sessions.json'), 'utf8')) as {
    sessions: Array<{ id: string; messages: ChatMessage[] }>
  }
  return stored.sessions.find((item) => item.id === sessionId)?.messages ?? []
}

async function waitForDiskAssistant(
  userData: string,
  sessionId: string,
  matches: (message: ChatMessage) => boolean
): Promise<ChatMessage> {
  for (let attempt = 0; attempt < 100; attempt++) {
    const message = (await readDiskMessages(userData, sessionId)).findLast((entry) => entry.role === 'assistant')
    if (message && matches(message)) return message
    await wait(10)
  }
  const messages = await readDiskMessages(userData, sessionId)
  assert.fail(`Timed out waiting for assistant checkpoint: ${JSON.stringify(messages.at(-1))}`)
}

function assistant(messages: ChatMessage[]): ChatMessage {
  const result = messages.findLast((message) => message.role === 'assistant')
  assert.ok(result, 'session should contain an assistant message')
  return result
}

async function main(): Promise<void> {
  const tempRoot = await mkdtemp(join(tmpdir(), 'ycode-chat-persistence-'))
  const state: PersistenceTestState = {
    userData: join(tempRoot, 'userData'),
    handlers: new Map(),
    promptCalls: 0,
    abortCalls: 0,
    cancelOutputEmitted: false,
    useOmp: true,
    fallbackHistory: null
  }
  await mkdir(state.userData, { recursive: true })

  const globals = globalThis as typeof globalThis & Record<string, unknown>
  const previousState = globals[TEST_STATE_KEY]
  globals[TEST_STATE_KEY] = state

  try {
    const ipcEntry = join(process.cwd(), 'src', 'main', 'ipc.ts')
    const sessionsEntry = join(process.cwd(), 'src', 'main', 'sessions.ts')
    const outfile = join(tempRoot, 'ipc-chat-persistence.mjs')
    await build({
      stdin: {
        contents: `export { registerIpc } from ${JSON.stringify(ipcEntry)};\nexport * as sessions from ${JSON.stringify(sessionsEntry)};`,
        resolveDir: process.cwd(),
        sourcefile: 'chat-persistence-regression-entry.ts',
        loader: 'ts'
      },
      outfile,
      bundle: true,
      external: ['node-pty'],
      platform: 'node',
      format: 'esm',
      banner: { js: "import { createRequire as createIpcTestRequire } from 'node:module'; const require = createIpcTestRequire(import.meta.url);" },
      logLevel: 'silent',
      plugins: [
        testTerminalHostPlugin,
        {
          name: 'chat-persistence-mocks',
          setup(pluginBuild) {
            pluginBuild.onResolve({ filter: /^electron$/ }, () => ({ path: 'electron', namespace: 'mock-electron' }))
            pluginBuild.onLoad({ filter: /.*/, namespace: 'mock-electron' }, () => ({
              loader: 'js',
              contents: `
                const state = globalThis.${TEST_STATE_KEY};
                export const app = {
                  getPath: (name) => name === 'userData' ? state.userData : state.userData,
                  getName: () => 'YCode',
                  getVersion: () => '0.1.0',
                  isPackaged: false, isReady: () => true,
                  getAppPath: () => state.userData,
                  getLoginItemSettings: () => ({ openAtLogin: false }),
                  setLoginItemSettings: () => {}
                };
                export const safeStorage = {
                  isEncryptionAvailable: () => false,
                  encryptString: (value) => Buffer.from(value),
                  decryptString: (value) => value.toString()
                };
                export const ipcMain = {
                  handle: (name, handler) => state.handlers.set(name, handler)
                };
                export const dialog = { showOpenDialog: async () => ({ canceled: true, filePaths: [] }) };
                export const shell = { openPath: async () => '' };
                export const powerSaveBlocker = { start: () => 1, stop: () => true, isStarted: () => false };
                export class BrowserWindow {}
                export class WebContentsView {}
              `
            }))
            pluginBuild.onResolve({ filter: /^\.\/omp$/ }, () => ({ path: 'omp', namespace: 'mock-omp' }))
            pluginBuild.onLoad({ filter: /.*/, namespace: 'mock-omp' }, () => ({
              loader: 'js',
              contents: `
                import { EventEmitter } from 'node:events';
                const state = globalThis.${TEST_STATE_KEY};
                const toolId = 'persisted-tool';
                const fullOutput = '${'long-output-'.repeat(60)}PERSISTED-OUTPUT-SENTINEL';
                const emit = (kernel, frame) => kernel.emit('event', frame);
                export class OmpKernel extends EventEmitter {
                  constructor() { super(); }
                  setThinkingLevel() { return Promise.resolve(); }
                  setHostTools() { return Promise.resolve({ type: 'response', success: true, data: { toolNames: [] } }); }
                  send() { return Promise.resolve({ type: 'response', success: true, data: { messages: [] } }); }
                  prompt(message) {
                    state.promptCalls++;
                    if (message.split('\\n', 1)[0] === 'persist-cancel') {
                      setImmediate(() => {
                        emit(this, { type: 'message_update', assistantMessageEvent: { type: 'text_delta', delta: 'cancel-before' } });
                        emit(this, { type: 'tool_execution_start', toolCall: { id: toolId, name: 'shell', args: { command: 'long-running' } } });
                        emit(this, { type: 'tool_execution_update', toolCall: { id: toolId, partialResult: 'cancelled tool output retained' } });
                        state.cancelOutputEmitted = true;
                      });
                    } else {
                      setImmediate(() => {
                        emit(this, { type: 'message_end', message: { role: 'assistant', usage: { input: 1200, output: 10, totalTokens: 1210 } } });
                        emit(this, { type: 'message_update', assistantMessageEvent: { type: 'text_delta', delta: 'before' } });
                        emit(this, { type: 'tool_execution_start', toolCall: { id: toolId, name: 'read_file', args: { path: 'large.txt' } } });
                        emit(this, { type: 'tool_execution_update', toolCall: { id: toolId, partialResult: fullOutput } });
                        emit(this, { type: 'message_update', assistantMessageEvent: { type: 'text_delta', delta: 'after' } });
                        emit(this, { type: 'tool_execution_end', toolCall: { id: toolId, name: 'read_file', result: { content: [{ type: 'text', text: fullOutput }], isError: false }, success: true } });
                        emit(this, { type: 'message_end', message: { role: 'assistant', usage: { contextTokens: 12800, input: 12000, cacheRead: 800, output: 200, totalTokens: 13000 } } });
                        emit(this, { type: 'agent_end', yielded: true });
                      });
                    }
                    return Promise.resolve({ type: 'response', success: true });
                  }
                  abort() {
                    state.abortCalls++;
                    setImmediate(() => emit(this, { type: 'agent_end', yielded: true }));
                    return Promise.resolve({ type: 'response', success: true });
                  }
                  dispose() {}
                }
                export const detectOmpPath = () => state.useOmp ? 'mock-omp' : null;
                export const getOmpInfo = () => ({ available: true, exePath: 'mock-omp', profile: 'ycode', model: 'ycode/test' });
              `
            }))
            pluginBuild.onResolve({ filter: /^\.\/agent\/loop$/ }, () => ({ path: 'agent-loop', namespace: 'mock-loop' }))
            pluginBuild.onLoad({ filter: /.*/, namespace: 'mock-loop' }, () => ({
              loader: 'js',
              contents: `
                const state = globalThis.${TEST_STATE_KEY};
                export async function runTurn(options) {
                  state.fallbackHistory = options.history;
                  state.fallbackWireHistory = options.wireHistory;
                  const content = 'fallback response';
                  options.emit({ runId: options.runId, type: 'text', delta: content });
                  options.emit({ runId: options.runId, type: 'done', content, tools: [] });
                  return { content, tools: [] };
                }
              `
            }))
            pluginBuild.onResolve({ filter: /^\.\/gateway$/ }, () => ({ path: 'gateway', namespace: 'mock-gateway' }))
            pluginBuild.onLoad({ filter: /.*/, namespace: 'mock-gateway' }, () => ({
              loader: 'js',
              contents: 'export class Gateway { constructor() {} }'
            }))
          }
        }
      ]
    })

    const ipc = await import(pathToFileURL(outfile).href)
    ;(ipc.registerIpc as (getWindow: () => unknown) => void)(() => ({
      webContents: { send: () => undefined }
    }))

    const invoke = (channel: string, ...args: unknown[]): Promise<unknown> => {
      const handler = state.handlers.get(channel)
      if (!handler) throw new Error(`Missing IPC handler: ${channel}`)
      return Promise.resolve(handler({}, ...args))
    }

    await invoke('models:save', { id: 'usage-snapshot-model', name: 'Usage snapshot model', providerId: 'custom',
      model: 'usage-test', baseUrl: 'http://127.0.0.1:1/v1', authMode: 'none', maxInputTokens: 64000 })
    await invoke('models:select', 'usage-snapshot-model')
    const normalSession = (await invoke('session:create', 'ordered persistence')) as { id: string }
    await invoke('chat:send', { sessionId: normalSession.id, text: 'persist-order' })
    const normalMessages = await waitForMessages(invoke, normalSession.id, (messages) => assistant(messages).status === 'done')
    const renamedSession = ((await invoke('session:list')) as Array<{ id: string; title: string }>)
      .find((entry) => entry.id === normalSession.id)
    assert.equal(renamedSession?.title, 'persist-order', 'the first send should set the session title')
    const normalAssistant = assistant(normalMessages)
    const expectedParts = [
      { type: 'text', text: 'before' },
      { type: 'tool', toolId: 'persisted-tool' },
      { type: 'text', text: 'after' }
    ]
    assert.deepEqual(normalAssistant.parts, expectedParts, 'session:messages should retain text/tool/text order')
    assert.equal(normalAssistant.status, 'done')
    assert.equal(normalAssistant.tools?.length, 1)
    const normalTool = normalAssistant.tools?.[0]
    assert.ok(normalTool, 'completed tool should be present')
    assert.ok((normalTool.output?.length ?? 0) > 600, 'full tool output should exceed the summary limit')
    assert.ok(normalTool.output?.endsWith('PERSISTED-OUTPUT-SENTINEL'), 'full tool output sentinel should be retained')
    assert.ok((normalTool.summary?.length ?? Infinity) < (normalTool.output?.length ?? 0),
      'the summary may be truncated while full output remains intact')

    const diskNormalAssistant = assistant(await readDiskMessages(state.userData, normalSession.id))
    assert.deepEqual(normalAssistant.contextUsage, {
      usage: { inputTokens: 12800, outputTokens: 200, totalTokens: 13000, cachedTokens: 800 }, contextWindow: 64000,
      modelId:'usage-snapshot-model',requestId:normalAssistant.requestId,source:'upstream',sampledAt:normalAssistant.contextUsage?.sampledAt
    }, 'latest tool-round input is a sample, not the sum of all round usage')
    assert.ok((normalAssistant.contextUsage?.sampledAt??0)>=normalAssistant.createdAt,'the persisted sample identifies its observation time')
    assert.deepEqual(diskNormalAssistant.contextUsage, normalAssistant.contextUsage,
      'the checkpoint whitelist must persist the upstream sample and frozen denominator together')
    assert.deepEqual(diskNormalAssistant.parts, expectedParts, 'disk JSON should retain text/tool/text order')
    assert.ok((diskNormalAssistant.tools?.[0].output?.length ?? 0) > 600,
      'disk JSON should retain output beyond the summary limit')
    assert.equal(diskNormalAssistant.tools?.[0].output, normalTool.output,
      'disk JSON should retain the complete tool output')
    assert.ok(diskNormalAssistant.tools?.[0].output?.endsWith('PERSISTED-OUTPUT-SENTINEL'))

    const fallbackSession = (await invoke('session:create', 'fallback history')) as { id: string }
    const sessions = ipc.sessions as {
      appendMessage(sessionId: string, message: ChatMessage): Promise<void>
    }
    await sessions.appendMessage(fallbackSession.id, {
      id: 'fallback-user',
      role: 'user',
      content: 'old request',
      quotes: [{ id: 'saved-table', sourceId: 'table-source', kind: 'table', markdown: '| name | value |\n| --- | --- |\n| A | 42 |', preview: 'A 42' }],
      createdAt: Date.now(),
      status: 'done'
    })
    await sessions.appendMessage(fallbackSession.id, {
      id: 'fallback-assistant',
      role: 'assistant',
      content: 'old answer',
      parts: [{ type: 'text', text: 'old answer' }],
      createdAt: Date.now(),
      status: 'done'
    })
    state.useOmp = false
    await invoke('chat:send', { sessionId: fallbackSession.id, text: 'fallback-history' })
    const fallbackMessages = await waitForMessages(
      invoke,
      fallbackSession.id,
      (messages) => assistant(messages).content === 'fallback response' && assistant(messages).status === 'done'
    )
    assert.equal(assistant(fallbackMessages).content, 'fallback response')
    assert.deepEqual(state.fallbackHistory?.map((message) => message.content), ['old request', 'old answer'],
      'the fallback loop should receive previously persisted user and assistant messages')
    assert.equal(state.fallbackWireHistory?.[0].content, 'old request\n\n| name | value |\n| --- | --- |\n| A | 42 |',
      'the next fallback request must include the exact previously quoted table snapshot')
    state.useOmp = true

    const cancelSession = (await invoke('session:create', 'cancelled persistence')) as { id: string }
    const cancelledSend = (await invoke('chat:send', {
      sessionId: cancelSession.id,
      text: 'persist-cancel'
    })) as { runId: string }
    for (let attempt = 0; attempt < 100 && !state.cancelOutputEmitted; attempt++) await wait(10)
    assert.equal(state.cancelOutputEmitted, true, 'fake OMP should emit a running tool and partial output')
    const streamingCheckpoint = await waitForDiskAssistant(state.userData, cancelSession.id, (message) =>
      message.status === 'streaming' &&
      message.content === 'cancel-before' &&
      message.tools?.[0].status === 'running' &&
      message.tools?.[0].output === 'cancelled tool output retained'
    )
    assert.deepEqual(streamingCheckpoint.parts, [
      { type: 'text', text: 'cancel-before' },
      { type: 'tool', toolId: 'persisted-tool' }
    ], 'the streaming checkpoint should persist text/tool order before cancellation')
    const cancelResult = (await invoke('chat:cancel', cancelledSend.runId)) as { cancelled: boolean }
    assert.equal(cancelResult.cancelled, true)

    const cancelledMessages = await waitForMessages(
      invoke,
      cancelSession.id,
      (messages) => assistant(messages).status === 'cancelled'
    )
    const cancelledAssistant = assistant(cancelledMessages)
    assert.deepEqual(cancelledAssistant.parts, [
      { type: 'text', text: 'cancel-before' },
      { type: 'tool', toolId: 'persisted-tool' }
    ])
    assert.equal(cancelledAssistant.status, 'cancelled')
    assert.equal(cancelledAssistant.tools?.[0].status, 'cancelled')
    assert.equal(cancelledAssistant.tools?.[0].output, 'cancelled tool output retained')

    const diskCancelledAssistant = assistant(await readDiskMessages(state.userData, cancelSession.id))
    assert.equal(diskCancelledAssistant.status, 'cancelled', 'disk JSON should persist the assistant cancellation')
    assert.equal(diskCancelledAssistant.tools?.[0].status, 'cancelled', 'disk JSON should persist terminal tool status')
    assert.equal(diskCancelledAssistant.tools?.[0].output, 'cancelled tool output retained',
      'disk JSON should preserve output from a running tool cancelled mid-turn')
    assert.equal(state.promptCalls, 2)
    assert.equal(state.abortCalls, 1)

    console.log('chat-persistence-regression: all checks passed')
  } finally {
    globals[TEST_STATE_KEY] = previousState
    await rm(tempRoot, { recursive: true, force: true })
  }
}

void main().catch((error: unknown) => {
  console.error('chat-persistence-regression failed:', error)
  process.exitCode = 1
})