import type { Message, UserMessage } from "@earendil-works/pi-ai/compat";
import {
  createFindToolDefinition,
  createGrepToolDefinition,
  createLsToolDefinition,
  createReadToolDefinition,
} from "@earendil-works/pi-coding-agent";
import type { ExtensionContext, ToolDefinition } from "@earendil-works/pi-coding-agent";
import type { BtwSettings } from "./config";
import { consumeSideStream, defaultSideModelClient, type SideProgressEvent, type SideRuntime } from "./side-runtime.ts";

const SIDE_SAFETY = `You are answering a SIDE QUESTION about the current coding session.
You MAY use the read-only tools (read, grep, find, ls) to investigate the codebase when the answer needs it.
Rules:
- Prefer answering directly from the conversation context above plus your general knowledge; use tools only when that is insufficient.
- Use tools ONLY to answer this side question. Make NO edits, run NO shell commands, and do not use write/edit tools (you do not have them).
- Do NOT continue, plan, or perform the user's main task.
- Use as few tool calls as possible.
- If you still cannot determine the answer, say so plainly.
Keep the answer concise.`;

export function buildSideQuestion(question: string, priorDigest?: string, now = Date.now): UserMessage {
  const prior = priorDigest ? `Earlier in this side thread:\n${priorDigest}\n\n` : "";
  return {
    role: "user",
    content: [{ type: "text", text: `${prior}${SIDE_SAFETY}\n\nSide question: ${question}` }],
    timestamp: now(),
  };
}

export type SideResult = { text: string; aborted: boolean; error?: string; toolsUsed: string[] };

export type SideRunDependencies = Partial<SideRuntime>;

const defaultToolFactory = (cwd: string): ToolDefinition<any, any>[] => [
  createReadToolDefinition(cwd),
  createGrepToolDefinition(cwd),
  createFindToolDefinition(cwd),
  createLsToolDefinition(cwd),
];

function resolveRuntime(deps: SideRunDependencies | undefined): SideRuntime {
  return {
    modelClient: deps?.modelClient ?? defaultSideModelClient,
    toolFactory: deps?.toolFactory ?? defaultToolFactory,
    now: deps?.now ?? Date.now,
    ...(deps?.onProgress ? { onProgress: deps.onProgress } : {}),
  };
}

type ToolResultMsg = Extract<Message, { role: "toolResult" }>;
function toolResultMsg(id: string, name: string, text: string, isError: boolean, timestamp: number): ToolResultMsg {
  return { role: "toolResult", toolCallId: id, toolName: name, content: [{ type: "text", text }], isError, timestamp };
}

export async function runSide(
  ctx: ExtensionContext,
  opts: { prefix: Message[]; tail?: string; question: string; settings: BtwSettings; signal: AbortSignal; onProgress?: (event: SideProgressEvent) => void },
  deps?: SideRunDependencies,
): Promise<SideResult> {
  const runtime = resolveRuntime(deps);
  const progress = (event: SideProgressEvent) => {
    try { runtime.onProgress?.(event); } catch {}
    try { opts.onProgress?.(event); } catch {}
  };
  const model = ctx.model;
  if (!model) return { text: "", aborted: false, error: "No model selected", toolsUsed: [] };

  const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
  if (!auth.ok) {
    return { text: "", aborted: false, toolsUsed: [], error: auth.error };
  }

  const cwd = (ctx as { cwd?: string }).cwd ?? process.cwd();
  const allDefs = runtime.toolFactory(cwd);
  const allowed = allDefs.filter((t) => opts.settings.toolAllowlist.includes(t.name));
  const defByName = new Map(allowed.map((t) => [t.name, t] as const));
  const tools = allowed.map((t) => ({ name: t.name, description: t.description, parameters: t.parameters }));
  const budget = opts.settings.toolCallBudget;

  const messages: Message[] = [...opts.prefix, buildSideQuestion(opts.question, opts.tail, runtime.now)];
  const toolsUsed = new Set<string>();
  let toolCallCount = 0;
  let iterations = 0;
  let request = 0;

  try {
    while (true) {
      if (opts.signal.aborted) return { text: "", aborted: true, toolsUsed: [...toolsUsed] };

      // Once the budget is spent, drop tools so the model is forced to answer.
      const budgetSpent = toolCallCount >= budget;
      const activeTools = budgetSpent ? [] : tools;

      const requestNumber = ++request;
      progress({ type: "model-request", request: requestNumber });
      const context = { systemPrompt: ctx.getSystemPrompt(), messages, tools: activeTools };
      const requestOptions = { apiKey: auth.apiKey, headers: auth.headers, env: auth.env, maxTokens: opts.settings.answerMaxTokens, cacheRetention: "short" as const, signal: opts.signal };
      // Select a capability before issuing the request. A broken stream never falls
      // back to complete, avoiding a duplicate model/tool execution.
      let response;
      if (runtime.modelClient.stream) {
        const outcome = await consumeSideStream(
          runtime.modelClient.stream(model, context, requestOptions),
          opts.signal,
          (text) => progress({ type: "text", request: requestNumber, text }),
        );
        const terminalMessage = outcome.kind === "done"
          ? outcome.message
          : outcome.message?.stopReason === "error" || outcome.message?.stopReason === "aborted"
            ? outcome.message
            : undefined;
        if (terminalMessage) {
          progress({ type: "model-response", request: requestNumber, stopReason: terminalMessage.stopReason });
        }
        if (outcome.kind === "aborted") return { text: "", aborted: true, toolsUsed: [...toolsUsed] };
        if (outcome.kind === "error") return { text: outcome.text, aborted: false, error: outcome.error, toolsUsed: [...toolsUsed] };
        response = outcome.message;
      } else if (runtime.modelClient.complete) {
        response = await runtime.modelClient.complete(model, context, requestOptions);
        progress({ type: "model-response", request: requestNumber, stopReason: response.stopReason });
      } else {
        return { text: "", aborted: false, error: "Side model client has no stream or complete capability", toolsUsed: [] };
      }
      if (opts.signal.aborted || response.stopReason === "aborted") return { text: "", aborted: true, toolsUsed: [...toolsUsed] };
      if (response.stopReason === "error") {
        return {
          text: response.content.filter((block): block is { type: "text"; text: string } => block.type === "text").map((block) => block.text).join("\n"),
          aborted: false,
          error: response.errorMessage || "provider error",
          toolsUsed: [...toolsUsed],
        };
      }
      messages.push(response);

      const calls = response.content.filter(
        (c): c is { type: "toolCall"; id: string; name: string; arguments: Record<string, unknown> } =>
          c.type === "toolCall",
      );

      if (calls.length === 0) {
        const text = response.content
          .filter((c): c is { type: "text"; text: string } => c.type === "text")
          .map((c) => c.text)
          .join("\n");
        const notice = budgetSpent ? `\n\n_(stopped after the ${budget}-tool-call budget)_` : "";
        return { text: text + notice, aborted: false, toolsUsed: [...toolsUsed] };
      }

      // Every toolCall in the batch needs a matching toolResult before the next complete().
      for (const call of calls) {
        const def = defByName.get(call.name);
        if (!def) {
          messages.push(toolResultMsg(call.id, call.name, `Tool "${call.name}" is not available.`, true, runtime.now()));
          continue;
        }
        if (toolCallCount >= budget) {
          messages.push(toolResultMsg(call.id, call.name, `Tool-call budget of ${budget} reached. Answer now from what you have.`, true, runtime.now()));
          continue;
        }
        toolCallCount++;
        toolsUsed.add(call.name);
        progress({ type: "tool-start", name: call.name, id: call.id });
        try {
          const r = await def.execute(call.id, call.arguments as never, opts.signal, undefined, ctx);
          messages.push({ role: "toolResult", toolCallId: call.id, toolName: call.name, content: r.content, isError: false, timestamp: runtime.now() });
          progress({ type: "tool-end", name: call.name, id: call.id, isError: false });
        } catch (e) {
          messages.push(toolResultMsg(call.id, call.name, e instanceof Error ? e.message : String(e), true, runtime.now()));
          progress({ type: "tool-end", name: call.name, id: call.id, isError: true });
        }
      }

      // Safety net against a pathological non-terminating model (should not trigger:
      // once budget is spent activeTools=[] forces a final text answer next pass).
      if (++iterations > budget + 3) {
        const text = messages
          .filter((m): m is Extract<Message, { role: "assistant" }> => m.role === "assistant")
          .flatMap((m) => m.content)
          .filter((c): c is { type: "text"; text: string } => c.type === "text")
          .map((c) => c.text)
          .join("\n");
        return { text: text || "Stopped: deep investigation exceeded its step limit.", aborted: false, toolsUsed: [...toolsUsed] };
      }
    }
  } catch (err) {
    if (opts.signal.aborted) return { text: "", aborted: true, toolsUsed: [...toolsUsed] };
    return { text: "", aborted: false, toolsUsed: [...toolsUsed], error: err instanceof Error ? err.message : String(err) };
  }
}

const REFINE_INSTRUCTION = `Rewrite the side Q/A below into a short note for the agent working on the main task.
Rules:
- Keep ONLY what helps the main task: the conclusion, key facts, decisions, and file/symbol references.
- Drop hedging, restated context, and exploration narrative.
- At most 6 lines. No preamble, no sign-off. Output the note body only.`;

export type QuickResult = { text: string; aborted: boolean; error?: string };

/** One quick pass that rewrites a shared Q/A into a main-actionable note. */
export async function runRefine(
  ctx: ExtensionContext,
  opts: { prefix: Message[]; question: string; answer: string; settings: BtwSettings; signal: AbortSignal },
  deps?: Pick<SideRunDependencies, "modelClient" | "now" | "onProgress">,
): Promise<QuickResult> {
  const runtime = resolveRuntime(deps);
  const model = ctx.model;
  if (!model) return { text: "", aborted: false, error: "No model selected" };

  const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
  if (!auth.ok) return { text: "", aborted: false, error: auth.error };

  const refineMsg: UserMessage = {
    role: "user",
    content: [{ type: "text", text: `${REFINE_INSTRUCTION}\n\nQ: ${opts.question}\nA: ${opts.answer}` }],
    timestamp: runtime.now(),
  };
  const messages: Message[] = [...opts.prefix, refineMsg];

  const progress = (event: SideProgressEvent) => {
    try {
      runtime.onProgress?.(event);
    } catch {
      // Progress observers are optional and cannot interrupt a billed request.
    }
  };

  if (opts.signal.aborted) return { text: "", aborted: true };

  try {
    progress({ type: "model-request", request: 1 });
    const context = { systemPrompt: ctx.getSystemPrompt(), messages, tools: [] };
    const requestOptions = { apiKey: auth.apiKey, headers: auth.headers, env: auth.env, maxTokens: opts.settings.refineMaxTokens, cacheRetention: "short" as const, signal: opts.signal };
    let response;
    if (runtime.modelClient.stream) {
      const outcome = await consumeSideStream(runtime.modelClient.stream(model, context, requestOptions), opts.signal);
      const terminalMessage = outcome.kind === "done"
        ? outcome.message
        : outcome.message?.stopReason === "error" || outcome.message?.stopReason === "aborted"
          ? outcome.message
          : undefined;
      if (terminalMessage) {
        progress({ type: "model-response", request: 1, stopReason: terminalMessage.stopReason });
      }
      if (outcome.kind === "aborted") return { text: "", aborted: true };
      if (outcome.kind === "error") return { text: outcome.text.trim(), aborted: false, error: outcome.error };
      response = outcome.message;
    } else if (runtime.modelClient.complete) {
      response = await runtime.modelClient.complete(model, context, requestOptions);
      progress({ type: "model-response", request: 1, stopReason: response.stopReason });
    } else {
      return { text: "", aborted: false, error: "Side model client has no stream or complete capability" };
    }
    if (opts.signal.aborted || response.stopReason === "aborted") return { text: "", aborted: true };
    const text = response.content
      .filter((block): block is { type: "text"; text: string } => block.type === "text")
      .map((block) => block.text)
      .join("\n");
    if (response.stopReason === "error") return { text: text.trim(), aborted: false, error: response.errorMessage || "provider error" };
    return { text: text.trim(), aborted: false };
  } catch (err) {
    if (opts.signal.aborted) return { text: "", aborted: true };
    return { text: "", aborted: false, error: err instanceof Error ? err.message : String(err) };
  }
}