import { WORKSPACE_STORAGE_TYPE } from '@learnitive/constants';
import type { ProjectAiToolCall } from '@learnitive/project-ai';
import { projectFilePath } from '@learnitive/project-storage';
import { createTestEnvironment } from '@learnitive/test-utils';
import { describe, expect, it } from 'vitest';
import type { ProjectAiProposal, ProjectAiWorkspaceState } from '../project-ai-service';

const WS_NAME = 'ai-tool-hardening-project';
const MAIN_TEX = `${WS_NAME}:main.tex`;

async function setup() {
  const env = createTestEnvironment();
  const services = env.instantiateAll();
  await env.mountAll();
  await services.workspaceOps.createWorkspaceInfo({
    name: WS_NAME,
    type: WORKSPACE_STORAGE_TYPE.Memory,
    metadata: {},
  });
  const source = '\\documentclass{article}\n\\begin{document}\nBefore after\n\\end{document}\n';
  await services.fileSystem.createTextFile(MAIN_TEX, source);
  await services.latexSource.load(MAIN_TEX);
  return { env, services, source };
}

type ProjectAiInternals = {
  prepareProposal(wsName: string, calls: readonly ProjectAiToolCall[]): Promise<ProjectAiProposal>;
  runAgent(wsName: string, signal: AbortSignal, extraInstruction?: string): Promise<void>;
  complete(messages: readonly { role: string; content: string }[], signal: AbortSignal): Promise<string>;
  runtime(wsName: string): { state: ProjectAiWorkspaceState };
};

describe('ProjectAiService tool hardening', () => {
  it('recovers an empty insert_table call from the active editor context', async () => {
    const { env, services, source } = await setup();
    const offset = source.indexOf(' after');
    services.projectAi.setEditorContext(WS_NAME, {
      path: projectFilePath('main.tex'),
      line: 3,
      column: 7,
      offset,
    });

    const internal = services.projectAi as unknown as ProjectAiInternals;
    const proposal = await internal.prepareProposal(WS_NAME, [{
      id: 'empty-table',
      name: 'insert_table',
      arguments: {},
    }]);

    const operation = proposal.operations[0];
    expect(operation?.kind).toBe('replace-file');
    if (operation?.kind === 'replace-file') {
      expect(operation.oldText).toBe(source);
      expect(operation.newText).toContain('\\begin{table}[htbp]');
      expect(operation.baseRevision).toBe(0);
    }
    expect(services.latexSource.getSnapshot(MAIN_TEX)?.text).toBe(source);
    env.controller.abort();
  });

  it('asks the model to repair an invalid mutation instead of failing the whole agent turn', async () => {
    const { env, services, source } = await setup();
    const internal = services.projectAi as unknown as ProjectAiInternals;
    const responses = [
      JSON.stringify({ message: '', toolCalls: [{ name: 'apply_patch', arguments: {} }] }),
      JSON.stringify({
        message: 'Prepared a reviewable correction.',
        toolCalls: [{
          name: 'apply_patch',
          arguments: { file: 'main.tex', replacements: [{ find: 'Before', replace: 'Changed' }] },
        }],
      }),
    ];
    const seenMessages: string[][] = [];
    internal.complete = async (messages) => {
      seenMessages.push(messages.map((message) => message.content));
      return responses.shift() ?? JSON.stringify({ message: 'done', toolCalls: [] });
    };

    await internal.runAgent(WS_NAME, new AbortController().signal);

    expect(seenMessages).toHaveLength(2);
    expect(seenMessages[1]?.some((message) => message.includes('TOOL_VALIDATION_ERROR'))).toBe(true);
    const proposal = services.projectAi.getWorkspaceState(WS_NAME).proposal;
    expect(proposal?.operations[0]?.kind).toBe('replace-file');
    expect(services.latexSource.getSnapshot(MAIN_TEX)?.text).toBe(source);
    env.controller.abort();
  });

  it('uses the generic project text path for active Python editor context', async () => {
    const { env, services } = await setup();
    const mainPy = `${WS_NAME}:main.py`;
    const pythonSource = 'def bubble_sort(values):\n    return values\n';
    await services.fileSystem.createTextFile(mainPy, pythonSource);
    services.projectAi.setEditorContext(WS_NAME, {
      path: projectFilePath('main.py'),
      line: 1,
      column: 1,
      offset: 0,
    });

    const internal = services.projectAi as unknown as ProjectAiInternals;
    const seenMessages: string[][] = [];
    internal.complete = async (messages) => {
      seenMessages.push(messages.map((message) => message.content));
      return JSON.stringify({ message: 'Ready to edit Python.', toolCalls: [] });
    };

    await expect(internal.runAgent(WS_NAME, new AbortController().signal)).resolves.toBeUndefined();
    expect(seenMessages[0]?.some((message) => message.includes('ACTIVE_SOURCE main.py:1:1'))).toBe(true);
    expect(seenMessages[0]?.some((message) => message.includes('def bubble_sort(values):'))).toBe(true);
    expect(services.projectAi.getWorkspaceState(WS_NAME).messages.at(-1)?.content).toBe('Ready to edit Python.');
    env.controller.abort();
  });

  it('repairs malformed internal protocol JSON without exposing it as assistant prose', async () => {
    const { env, services } = await setup();
    const internal = services.projectAi as unknown as ProjectAiInternals;
    const malformed = '{"message":"Abstract written","toolCalls":[{"name":"insert_equation","arguments":{"latex":"\\\\[x';
    const responses = [
      malformed,
      JSON.stringify({ message: 'Recovered safely.', toolCalls: [] }),
    ];
    const seenMessages: string[][] = [];
    internal.complete = async (messages) => {
      seenMessages.push(messages.map((message) => message.content));
      return responses.shift() ?? JSON.stringify({ message: 'Recovered safely.', toolCalls: [] });
    };

    await internal.runAgent(WS_NAME, new AbortController().signal);

    expect(seenMessages).toHaveLength(2);
    expect(seenMessages[1]?.some((message) => message.includes('PROTOCOL_ERROR'))).toBe(true);
    const visible = services.projectAi.getWorkspaceState(WS_NAME).messages.filter((message) => message.role === 'assistant');
    expect(visible.some((message) => message.content === malformed)).toBe(false);
    expect(visible.at(-1)?.content).toBe('Recovered safely.');
    env.controller.abort();
  });

  it('executes read tools before considering mutations returned in the same model response', async () => {
    const { env, services, source } = await setup();
    const internal = services.projectAi as unknown as ProjectAiInternals;
    const responses = [
      JSON.stringify({
        message: '',
        toolCalls: [
          { name: 'read_file', arguments: { file: 'main.tex' } },
          {
            name: 'apply_patch',
            arguments: { file: 'main.tex', replacements: [{ find: 'Before', replace: 'Premature' }] },
          },
        ],
      }),
      JSON.stringify({
        message: 'Prepared after reading the source.',
        toolCalls: [{
          name: 'apply_patch',
          arguments: { file: 'main.tex', replacements: [{ find: 'Before', replace: 'Reviewed' }] },
        }],
      }),
    ];
    const seenMessages: string[][] = [];
    internal.complete = async (messages) => {
      seenMessages.push(messages.map((message) => message.content));
      return responses.shift() ?? JSON.stringify({ message: 'done', toolCalls: [] });
    };

    await internal.runAgent(WS_NAME, new AbortController().signal);

    expect(seenMessages).toHaveLength(2);
    expect(seenMessages[1]?.some((message) => message.includes('MUTATION_DEFERRED'))).toBe(true);
    expect(seenMessages[1]?.some((message) => message.includes('TOOL_RESULT') && message.includes('read_file'))).toBe(true);
    const operation = services.projectAi.getWorkspaceState(WS_NAME).proposal?.operations[0];
    expect(operation?.kind).toBe('replace-file');
    if (operation?.kind === 'replace-file') {
      expect(operation.newText).toContain('Reviewed');
      expect(operation.newText).not.toContain('Premature');
    }
    expect(services.latexSource.getSnapshot(MAIN_TEX)?.text).toBe(source);
    env.controller.abort();
  });

  it('rejects a proposal when the source revision changed even if the text was later reverted', async () => {
    const { env, services, source } = await setup();
    const internal = services.projectAi as unknown as ProjectAiInternals;
    const proposal = await internal.prepareProposal(WS_NAME, [{
      id: 'patch-1',
      name: 'apply_patch',
      arguments: {
        file: 'main.tex',
        replacements: [{ find: 'Before', replace: 'Changed' }],
      },
    }]);

    services.latexSource.update(MAIN_TEX, source.replace('Before', 'Temporary'), 'source');
    services.latexSource.update(MAIN_TEX, source, 'source');
    expect(services.latexSource.getSnapshot(MAIN_TEX)?.revision).toBeGreaterThan(0);
    expect(services.latexSource.getSnapshot(MAIN_TEX)?.text).toBe(source);

    const runtime = internal.runtime(WS_NAME);
    runtime.state = { ...runtime.state, proposal };

    await expect(services.projectAi.applyProposal(WS_NAME)).rejects.toThrow(/stale/u);
    expect(services.latexSource.getSnapshot(MAIN_TEX)?.text).toBe(source);
    env.controller.abort();
  });
});

describe('ProjectAiService bounded compile tool', () => {
  it('executes compile_project without exposing arbitrary sandbox runtime tools', async () => {
    const { env, services } = await setup();
    const internal = services.projectAi as unknown as ProjectAiInternals;
    const compiledFiles: Array<string | undefined> = [];
    const seenMessages: string[][] = [];
    const responses = [
      JSON.stringify({ message: '', toolCalls: [{ id: 'compile-1', name: 'compile_project', arguments: { file: 'main.py' } }] }),
      JSON.stringify({ message: 'main.py compiled successfully.', toolCalls: [] }),
    ];
    services.projectAi.setCompileToolExecutor({
      compileProject: async (file) => {
        compiledFiles.push(file);
        return { status: 'success', exitCode: 0, stdout: 'ok', stderr: '', diagnostics: [], executionReference: 'exec-1', provider: 'python', runtime: 'python3' };
      },
    });
    internal.complete = async (messages) => {
      seenMessages.push(messages.map((message) => message.content));
      return responses.shift() ?? JSON.stringify({ message: 'done', toolCalls: [] });
    };

    await internal.runAgent(WS_NAME, new AbortController().signal);

    expect(compiledFiles).toEqual(['main.py']);
    expect(seenMessages[1]?.some((message) => message.includes('TOOL_RESULT') && message.includes('compile_project') && message.includes('"provider":"python"'))).toBe(true);
    expect(seenMessages.flat().some((message) => message.includes('RUNTIME_TOOLS_UNAVAILABLE'))).toBe(false);
    expect(services.projectAi.getWorkspaceState(WS_NAME).messages.at(-1)?.content).toBe('main.py compiled successfully.');
    env.controller.abort();
  });

  it('keeps run_command unavailable when only the bounded compile executor is attached', async () => {
    const { env, services } = await setup();
    const internal = services.projectAi as unknown as ProjectAiInternals;
    const seenMessages: string[][] = [];
    const responses = [
      JSON.stringify({ message: '', toolCalls: [{ id: 'runtime-1', name: 'run_command', arguments: { command: 'python3 main.py' } }] }),
      JSON.stringify({ message: 'I cannot run arbitrary commands in Ask mode.', toolCalls: [] }),
    ];
    services.projectAi.setCompileToolExecutor({ compileProject: async () => ({ status: 'success' }) });
    internal.complete = async (messages) => {
      seenMessages.push(messages.map((message) => message.content));
      return responses.shift() ?? JSON.stringify({ message: 'done', toolCalls: [] });
    };

    await internal.runAgent(WS_NAME, new AbortController().signal);

    expect(seenMessages[1]?.some((message) => message.includes('RUNTIME_TOOLS_UNAVAILABLE'))).toBe(true);
    env.controller.abort();
  });
});