import { Agent, type AgentTool, type StreamFn } from "@earendil-works/pi-agent-core";
import type { AssistantMessage, Model, SimpleStreamOptions } from "@earendil-works/pi-ai";
import { convertToLlm, createReadOnlyTools, type ModelRegistry } from "@earendil-works/pi-coding-agent";
import { createReadChangeTool } from "./change-tool.js";
import { validateProjectPath } from "./path-guard.js";
import { buildCoverageCorrectionPrompt } from "./prompts.js";
import { addUsage, emptyReviewUsage, type ChangeSnapshot, type ReviewAgentResult } from "./types.js";

const MIN_INITIAL_TURNS = 20;
const MAX_INITIAL_TURNS = 40;
const INITIAL_ANALYSIS_RESERVE = 8;
const MIN_CORRECTION_TURNS = 8;
const MAX_CORRECTION_TURNS = 20;
const CORRECTION_ANALYSIS_RESERVE = 4;
const FINALIZATION_RESERVE_TURNS = 2;
const MAX_IDENTICAL_TOOL_BATCHES = 3;
export const MAX_REVIEW_TURNS = MAX_INITIAL_TURNS + MAX_CORRECTION_TURNS;
export const REVIEW_TOOL_NAMES = ["read_change", "read", "grep", "find", "ls"] as const;
const ALLOWED_TOOLS = new Set<string>(REVIEW_TOOL_NAMES);

export type ReviewPhase = "initial" | "coverage";
type PhaseUpdate = (message: string) => void;

interface ReviewSafetyFailure {
  kind: "phase-limit" | "global-limit" | "repeated-tool-batch";
  phase: ReviewPhase;
  phaseTurns: number;
  phaseLimit: number;
  totalTurns: number;
  missingPaths: string[];
  repeatedBatches?: number;
}

export function reviewPhaseTurnLimit(phase: ReviewPhase, evidencePages: number): number {
  const pages = Math.max(0, Math.floor(evidencePages));
  if (phase === "initial") {
    return Math.max(MIN_INITIAL_TURNS, Math.min(MAX_INITIAL_TURNS, pages + INITIAL_ANALYSIS_RESERVE));
  }
  return Math.max(MIN_CORRECTION_TURNS, Math.min(MAX_CORRECTION_TURNS, pages + CORRECTION_ANALYSIS_RESERVE));
}

function phaseLabel(phase: ReviewPhase): string {
  return phase === "initial" ? "initial review" : "coverage correction";
}

function stableSerialize(value: unknown): string {
  if (value === null) return "null";
  if (Array.isArray(value)) return `[${value.map(stableSerialize).join(",")}]`;
  if (typeof value === "object") {
    const object = value as Record<string, unknown>;
    return `{${Object.keys(object)
      .sort()
      .map((key) => `${JSON.stringify(key)}:${stableSerialize(object[key])}`)
      .join(",")}}`;
  }
  if (typeof value === "bigint") return `${value.toString()}n`;
  return JSON.stringify(value) ?? String(value);
}

function missingEvidenceText(paths: string[]): string {
  if (paths.length === 0) return "Immutable evidence coverage was complete.";
  const shown = paths.slice(0, 5).map((path) => JSON.stringify(path));
  const omitted = paths.length - shown.length;
  return `Missing immutable evidence for ${paths.length} path${paths.length === 1 ? "" : "s"}: ${shown.join(", ")}${omitted > 0 ? `, and ${omitted} more` : ""}.`;
}

function safetyFailureMessage(failure: ReviewSafetyFailure): string {
  const phase = phaseLabel(failure.phase);
  const progress = `${failure.phaseTurns}/${failure.phaseLimit} phase turns; ${failure.totalTurns}/${MAX_REVIEW_TURNS} total turns`;
  const coverage = missingEvidenceText(failure.missingPaths);
  const guidance = "Batch independent tool calls, retry with a model that supports parallel tool use, or split the changes.";

  if (failure.kind === "repeated-tool-batch") {
    return `Reviewer stopped during ${phase} after repeating the same tool batch ${failure.repeatedBatches} times without new immutable-evidence coverage (${progress}). ${coverage} ${guidance}`;
  }
  if (failure.kind === "global-limit") {
    return `Reviewer exceeded the ${MAX_REVIEW_TURNS}-turn global safety limit during ${phase} (${progress}). ${coverage} ${guidance}`;
  }
  return `Reviewer exceeded its ${phase} turn budget (${progress}). ${coverage} ${guidance}`;
}

function finalizationPrompt(phase: ReviewPhase, phaseTurns: number, phaseLimit: number, missingPaths: string[]): string {
  const remaining = Math.max(1, phaseLimit - phaseTurns);
  const coverage =
    missingPaths.length === 0
      ? "Immutable evidence coverage is complete."
      : `Immutable evidence remains incomplete for ${missingPaths.length} path${missingPaths.length === 1 ? "" : "s"}.`;
  if (remaining === 1) {
    return `Safety budget notice for ${phaseLabel(phase)}: one assistant turn remains. ${coverage} Stop using tools and return the entire final report now.`;
  }
  return `Safety budget notice for ${phaseLabel(phase)}: ${remaining} assistant turns remain. ${coverage} In the next turn, batch every essential independent tool call in parallel. Use the final turn without tools to return the entire report in the required format.`;
}

function outputLimit(model: Model<any>): number {
  return Math.max(256, Math.min(8192, model.maxTokens, Math.floor(model.contextWindow * 0.15)));
}

export function assertReviewFitsContext(
  model: Model<any>,
  snapshot: ChangeSnapshot,
  systemPrompt: string,
  prompt: string,
): number {
  const maxOutput = outputLimit(model);
  const promptTokens = Math.ceil(Buffer.byteLength(`${systemPrompt}\n${prompt}`) / 3);
  const explorationReserve = Math.max(4_000, Math.floor(model.contextWindow * 0.1));
  const required = snapshot.estimatedTokens + promptTokens + maxOutput + explorationReserve;
  if (required > model.contextWindow) {
    throw new Error(
      `${model.provider}/${model.id} has a ${model.contextWindow.toLocaleString()} token context window, but a complete review is estimated to require ${required.toLocaleString()} tokens. Choose a larger-context model or split the changes.`,
    );
  }
  return maxOutput;
}

function lastAssistant(messages: readonly unknown[]): AssistantMessage | undefined {
  for (let i = messages.length - 1; i >= 0; i -= 1) {
    const message = messages[i] as { role?: string };
    if (message.role === "assistant") return messages[i] as AssistantMessage;
  }
  return undefined;
}

function assistantText(message: AssistantMessage): string {
  return message.content
    .filter((part): part is { type: "text"; text: string } => part.type === "text")
    .map((part) => part.text)
    .join("\n")
    .trim();
}

function safeOptions(
  options: SimpleStreamOptions | undefined,
  auth: {
    apiKey?: string;
    headers?: Record<string, string>;
    env?: Record<string, string>;
  },
  maxTokens: number,
  reasoning: boolean,
): SimpleStreamOptions {
  const result: SimpleStreamOptions = { ...options, maxTokens };
  if (reasoning) result.reasoning = "high";
  else delete result.reasoning;
  if (auth.apiKey !== undefined) result.apiKey = auth.apiKey;
  if (auth.headers !== undefined) result.headers = auth.headers;
  if (auth.env !== undefined) result.env = auth.env;
  return result;
}

export interface RunReviewAgentOptions {
  model: Model<any>;
  modelRegistry: ModelRegistry;
  snapshot: ChangeSnapshot;
  systemPrompt: string;
  prompt: string;
  signal?: AbortSignal;
  onPhase?: PhaseUpdate;
}

export async function runReviewAgent(options: RunReviewAgentOptions): Promise<ReviewAgentResult> {
  const { model, modelRegistry, snapshot, systemPrompt, prompt, signal, onPhase } = options;
  const maxTokens = assertReviewFitsContext(model, snapshot, systemPrompt, prompt);
  const provider = modelRegistry.getProvider(model.provider);
  if (!provider) throw new Error(`Provider ${JSON.stringify(model.provider)} is not available.`);
  const auth = await modelRegistry.getApiKeyAndHeaders(model);
  if (!auth.ok) throw new Error(auth.error);

  const { tool: readChange, coverage } = createReadChangeTool(snapshot);
  const repositoryTools = createReadOnlyTools(snapshot.root);
  const tools: AgentTool[] = [readChange, ...repositoryTools];
  const streamFn: StreamFn = (_requestedModel, context, streamOptions) =>
    provider.streamSimple(model, context, safeOptions(streamOptions, auth, maxTokens, model.reasoning));

  let phase: ReviewPhase = "initial";
  let phaseTurns = 0;
  let phaseLimit = reviewPhaseTurnLimit(phase, coverage.missingPageCount());
  let totalTurns = 0;
  let safetyFailure: ReviewSafetyFailure | undefined;
  let finalizationWarningSent = false;
  let coverageAtTurnStart = coverage.coveredCharacters();
  let currentToolBatch: string[] = [];
  let previousToolBatch: string | undefined;
  let identicalToolBatches = 0;

  const agent = new Agent({
    initialState: {
      systemPrompt,
      model,
      thinkingLevel: model.reasoning ? "high" : "off",
      tools,
      messages: [],
    },
    convertToLlm,
    streamFn,
    toolExecution: "parallel",
    beforeToolCall: async ({ toolCall, args }) => {
      if (!ALLOWED_TOOLS.has(toolCall.name)) {
        return { block: true, reason: `${toolCall.name} is not allowed in read-only review mode.` };
      }
      if (toolCall.name === "read_change") return undefined;

      const path = args && typeof args === "object" && "path" in args ? (args as { path?: unknown }).path : undefined;
      try {
        await validateProjectPath(snapshot.root, typeof path === "string" ? path : undefined);
        return undefined;
      } catch (error) {
        return { block: true, reason: error instanceof Error ? error.message : String(error) };
      }
    },
  });

  const beginPhase = (nextPhase: ReviewPhase, evidencePages: number) => {
    phase = nextPhase;
    phaseTurns = 0;
    phaseLimit = reviewPhaseTurnLimit(nextPhase, evidencePages);
    finalizationWarningSent = false;
    coverageAtTurnStart = coverage.coveredCharacters();
    currentToolBatch = [];
    previousToolBatch = undefined;
    identicalToolBatches = 0;
  };

  const throwIfSafetyFailure = () => {
    if (safetyFailure) throw new Error(safetyFailureMessage(safetyFailure));
  };

  agent.subscribe((event) => {
    if (safetyFailure) return;

    if (event.type === "turn_start") {
      coverageAtTurnStart = coverage.coveredCharacters();
      currentToolBatch = [];
      return;
    }

    if (event.type === "tool_execution_start") {
      currentToolBatch.push(`${event.toolName}:${stableSerialize(event.args)}`);
      const args = event.args as { path?: unknown };
      const suffix = typeof args?.path === "string" ? ` ${args.path}` : "";
      onPhase?.(`${toolLabel(event.toolName)}${suffix}`);
      return;
    }

    if (event.type !== "turn_end") return;
    phaseTurns += 1;
    totalTurns += 1;
    if (
      event.message.role !== "assistant" ||
      !event.message.content.some((part) => part.type === "toolCall")
    ) {
      return;
    }

    const coverageAdvanced = coverage.coveredCharacters() > coverageAtTurnStart;
    const toolBatch = currentToolBatch.length > 0 ? currentToolBatch.slice().sort().join("\n") : undefined;
    if (toolBatch) {
      identicalToolBatches = toolBatch === previousToolBatch && !coverageAdvanced ? identicalToolBatches + 1 : 1;
      previousToolBatch = toolBatch;
    } else {
      identicalToolBatches = 0;
      previousToolBatch = undefined;
    }

    if (identicalToolBatches >= MAX_IDENTICAL_TOOL_BATCHES) {
      safetyFailure = {
        kind: "repeated-tool-batch",
        phase,
        phaseTurns,
        phaseLimit,
        totalTurns,
        missingPaths: coverage.missingPaths(),
        repeatedBatches: identicalToolBatches,
      };
      agent.abort();
      return;
    }

    if (totalTurns >= MAX_REVIEW_TURNS || phaseTurns >= phaseLimit) {
      safetyFailure = {
        kind: totalTurns >= MAX_REVIEW_TURNS ? "global-limit" : "phase-limit",
        phase,
        phaseTurns,
        phaseLimit,
        totalTurns,
        missingPaths: coverage.missingPaths(),
      };
      agent.abort();
      return;
    }

    if (!finalizationWarningSent && phaseTurns >= phaseLimit - FINALIZATION_RESERVE_TURNS) {
      finalizationWarningSent = true;
      onPhase?.("turn budget nearly exhausted; requesting completion");
      agent.steer({
        role: "user",
        content: [{ type: "text", text: finalizationPrompt(phase, phaseTurns, phaseLimit, coverage.missingPaths()) }],
        timestamp: Date.now(),
      });
    }
  });

  const abortAgent = () => agent.abort();
  if (signal?.aborted) throw new Error("Review cancelled.");
  signal?.addEventListener("abort", abortAgent, { once: true });

  try {
    onPhase?.("analyzing captured changes");
    await agent.prompt(prompt);
    throwIfSafetyFailure();

    let missing = coverage.missingPaths();
    if (missing.length > 0 && !signal?.aborted) {
      beginPhase("coverage", coverage.missingPageCount());
      onPhase?.(`checking coverage for ${missing.length} path${missing.length === 1 ? "" : "s"}`);
      await agent.prompt(buildCoverageCorrectionPrompt(missing));
      throwIfSafetyFailure();
      missing = coverage.missingPaths();
    }

    const final = lastAssistant(agent.state.messages);
    if (!final) throw new Error("Reviewer returned no assistant response.");
    if (final.stopReason === "error") throw new Error(final.errorMessage || "Reviewer model failed.");
    if (final.stopReason === "aborted" || signal?.aborted) {
      throwIfSafetyFailure();
      throw new Error("Review cancelled.");
    }

    const text = assistantText(final);
    if (!text) throw new Error("Reviewer returned no textual report.");

    const usage = emptyReviewUsage();
    for (const message of agent.state.messages) {
      if (message.role === "assistant") addUsage(usage, message.usage);
    }

    return { text, usage, missingCoverage: missing, stopReason: final.stopReason };
  } finally {
    signal?.removeEventListener("abort", abortAgent);
    agent.abort();
  }
}

function toolLabel(toolName: string): string {
  switch (toolName) {
    case "read_change":
      return "reading change";
    case "read":
      return "reading context";
    case "grep":
      return "searching code";
    case "find":
      return "finding files";
    case "ls":
      return "listing files";
    default:
      return toolName;
  }
}