import { createAgentSession, SessionManager, type AgentSessionEvent, type ExtensionContext } from "@earendil-works/pi-coding-agent";
import { envelopePrompt } from "../context/envelope.js";
import { parseAgentResult } from "../context/result-contracts.js";
import { PATH_RE } from "../security/privacy.js";
import { SCOPED_READ_TOOLS, scopedReadOnlyToolDefinitions } from "../security/worker-scope.js";
import type { AgentResult, TaskEnvelope, ThinkingLevel } from "../types.js";
import { minimalResourceLoader } from "./minimal-resource-loader.js";

export interface ProviderFailureClassification {
  status: 429 | 500;
  errorClass: "rate-limited" | "provider-5xx";
}

export type MicroAgentErrorClass = ProviderFailureClassification["errorClass"] | "aborted" | "authentication" | "timeout" | "model-unavailable" | "result-contract" | "micro-agent-failed";

interface MicroAgentCallUsage {
  input: number;
  output: number;
  cacheRead: number;
  totalTokens: number;
  cost: number;
  model: string;
  provider: string;
  thinkingLevel: ThinkingLevel;
  turnElapsedMs: number;
  turnIndex: number;
  callIndex: number;
}

export type MicroAgentEvent =
  | { type: "model.call.started"; model: string; provider: string; thinkingLevel: ThinkingLevel; turnIndex: number; callIndex: number }
  | (MicroAgentCallUsage & { type: "model.call.completed" })
  | (MicroAgentCallUsage & { type: "model.call.failed"; errorClass: MicroAgentErrorClass })
  | { type: "tool.started"; toolCallId: string; toolName: string }
  | { type: "tool.completed" | "tool.failed"; toolCallId: string; toolName: string; toolElapsedMs: number };

const EXPLICIT_PROVIDER_5XX_RE = /\bprovider-5xx\b/i;
const EXPLICIT_RATE_LIMITED_RE = /\brate-limited\b/i;
const ABORTED_RE = /\baborted\b/i;
const AUTHENTICATION_RE = /(no api key|no credentials|unauthori[sz]ed|forbidden|authentication|\b401\b|\b403\b)/i;
const PROVIDER_5XX_PHRASE_RE = /(internal server|service unavailable|bad gateway|gateway timeout)/i;
const TIMEOUT_RE = /(timed out|timeout)/i;
const MODEL_UNAVAILABLE_RE = /(no models match|unknown model|model.*(?:not found|unavailable))/i;
const RATE_LIMITED_RE = /(rate.?limit|too many requests|\b429\b|overloaded)/i;
const PROVIDER_5XX_NUMERIC_RE = /\b(?:HTTP\s*)?5\d\d\b(?!\s*(?:ms|msec|s|sec|secs|seconds?)\b)/i;
const RESULT_CONTRACT_RE = /(agentresult|result contract|result.*schema|assistant result|duplicate fact|output contract)/i;

function failureText(error: unknown): string {
  const record = error && typeof error === "object" ? error as Record<string, unknown> : undefined;
  const raw = typeof error === "string" ? error : error instanceof Error ? error.message : typeof record?.message === "string" ? record.message : "";
  return raw.replace(PATH_RE, " ");
}

function structuredStatus(error: unknown): number | undefined {
  const record = error && typeof error === "object" ? error as Record<string, unknown> : undefined;
  return [record?.status, record?.statusCode, record?.response && typeof record.response === "object" ? (record.response as Record<string, unknown>).status : undefined].find((value): value is number => typeof value === "number");
}

function statusErrorClass(status: number): MicroAgentErrorClass | undefined {
  if (status === 429) return "rate-limited";
  if (status === 401 || status === 403) return "authentication";
  if (status >= 500 && status < 600) return "provider-5xx";
  return undefined;
}

export function classifyMicroAgentFailure(error: unknown): MicroAgentErrorClass {
  const status = structuredStatus(error);
  const fromStatus = status === undefined ? undefined : statusErrorClass(status);
  if (fromStatus) return fromStatus;
  const text = failureText(error);
  if (EXPLICIT_PROVIDER_5XX_RE.test(text)) return "provider-5xx";
  if (EXPLICIT_RATE_LIMITED_RE.test(text)) return "rate-limited";
  if (ABORTED_RE.test(text)) return "aborted";
  if (AUTHENTICATION_RE.test(text)) return "authentication";
  if (PROVIDER_5XX_PHRASE_RE.test(text)) return "provider-5xx";
  if (TIMEOUT_RE.test(text)) return "timeout";
  if (MODEL_UNAVAILABLE_RE.test(text)) return "model-unavailable";
  if (RATE_LIMITED_RE.test(text)) return "rate-limited";
  if (PROVIDER_5XX_NUMERIC_RE.test(text)) return "provider-5xx";
  if (error instanceof SyntaxError || RESULT_CONTRACT_RE.test(text)) return "result-contract";
  return "micro-agent-failed";
}

export function classifyProviderFailure(error: unknown): ProviderFailureClassification | undefined {
  const errorClass = classifyMicroAgentFailure(error);
  if (errorClass === "rate-limited") return { status: 429, errorClass };
  if (errorClass === "provider-5xx") return { status: 500, errorClass };
  return undefined;
}

export function createMicroAgentEventObserver(modelId: string, thinkingLevel: ThinkingLevel, onEvent: ((event: MicroAgentEvent) => Promise<void>) | undefined, now = Date.now) {
  let turnIndex = -1;
  let callIndex = 0;
  let samples = 0;
  const sampledUsage: Array<{ input: number; output: number; cacheRead: number; totalTokens: number; cost: number }> = [];
  let activeCall: { startedAt: number; turnIndex: number; callIndex: number } | undefined;
  const toolTimers = new Map<string, { toolName: string; startedAt: number }>();
  let callbacks: Promise<void> = Promise.resolve();
  let callbackFailed = false;
  let callbackError: unknown;
  const emit = (event: MicroAgentEvent) => {
    if (!onEvent) return;
    callbacks = callbacks.then(() => onEvent(event)).catch((error) => {
      if (!callbackFailed) { callbackFailed = true; callbackError = error; }
    });
  };
  const failPrompt = (error: unknown, usage: { input: number; output: number; cacheRead: number; totalTokens: number; cost: number }, samplesBefore: number) => {
    const sample = activeCall ?? (samples === samplesBefore ? { startedAt: now(), turnIndex: Math.max(0, turnIndex), callIndex: ++callIndex } : undefined);
    if (!sample) return;
    const accounted = sampledUsage.slice(samplesBefore).reduce((total, entry) => ({ input: total.input + entry.input, output: total.output + entry.output, cacheRead: total.cacheRead + entry.cacheRead, totalTokens: total.totalTokens + entry.totalTokens, cost: total.cost + entry.cost }), { input: 0, output: 0, cacheRead: 0, totalTokens: 0, cost: 0 });
    const remaining = { input: Math.max(0, usage.input - accounted.input), output: Math.max(0, usage.output - accounted.output), cacheRead: Math.max(0, usage.cacheRead - accounted.cacheRead), totalTokens: Math.max(0, usage.totalTokens - accounted.totalTokens), cost: Math.max(0, usage.cost - accounted.cost) };
    activeCall = undefined;
    samples += 1;
    sampledUsage.push(remaining);
    const provider = modelId.split("/", 1)[0] ?? "unknown";
    emit({ type: "model.call.failed", model: modelId, provider, thinkingLevel, ...remaining, turnElapsedMs: Math.max(0, now() - sample.startedAt), turnIndex: sample.turnIndex, callIndex: sample.callIndex, errorClass: classifyMicroAgentFailure(error) });
  };
  return {
    observe(event: AgentSessionEvent): void {
      if (event.type === "agent_start") {
        turnIndex = -1;
        activeCall = undefined;
        toolTimers.clear();
        return;
      }
      if (event.type === "turn_start") {
        turnIndex += 1;
        const sample = { startedAt: now(), turnIndex, callIndex: ++callIndex };
        activeCall = sample;
        emit({ type: "model.call.started", model: modelId, provider: modelId.split("/", 1)[0] ?? "unknown", thinkingLevel, turnIndex, callIndex: sample.callIndex });
        return;
      }
      if (event.type === "message_end" && event.message.role === "assistant") {
        if (!activeCall) return;
        const sample = activeCall;
        activeCall = undefined;
        const provider = event.message.provider;
        const responseModel = event.message.responseModel ?? event.message.model;
        const failed = event.message.stopReason === "error" || event.message.stopReason === "aborted";
        const usage = { model: `${provider}/${responseModel}`, provider, thinkingLevel, input: event.message.usage.input, output: event.message.usage.output, cacheRead: event.message.usage.cacheRead, totalTokens: event.message.usage.totalTokens, cost: event.message.usage.cost.total, turnElapsedMs: Math.max(0, now() - sample.startedAt), turnIndex: sample.turnIndex, callIndex: sample.callIndex };
        samples += 1;
        sampledUsage.push({ input: usage.input, output: usage.output, cacheRead: usage.cacheRead, totalTokens: usage.totalTokens, cost: usage.cost });
        emit(failed ? { type: "model.call.failed", ...usage, errorClass: classifyMicroAgentFailure(event.message.errorMessage ?? event.message.stopReason) } : { type: "model.call.completed", ...usage });
        return;
      }
      if (event.type === "tool_execution_start") {
        toolTimers.set(event.toolCallId, { toolName: event.toolName, startedAt: now() });
        emit({ type: "tool.started", toolCallId: event.toolCallId, toolName: event.toolName });
        return;
      }
      if (event.type === "tool_execution_end") {
        const timer = toolTimers.get(event.toolCallId);
        toolTimers.delete(event.toolCallId);
        if (!timer) return;
        emit({ type: event.isError ? "tool.failed" : "tool.completed", toolCallId: event.toolCallId, toolName: timer.toolName, toolElapsedMs: Math.max(0, now() - timer.startedAt) });
      }
    },
    sampleCount: () => samples,
    failPrompt,
    async flush(): Promise<void> {
      await callbacks;
      if (callbackFailed) throw callbackError;
    },
  };
}

function lastAssistantText(messages: readonly unknown[]): string {
  const message = [...messages].reverse().find((entry) => typeof entry === "object" && entry !== null && (entry as { role?: string }).role === "assistant") as { content?: Array<{ type?: string; text?: string }> } | undefined;
  const text = message?.content?.filter((part) => part.type === "text").map((part) => part.text ?? "").join("\n");
  if (!text) throw new Error("Worker returned no assistant result");
  return text;
}

export class MicroAgentBackend {
  constructor(private readonly context: ExtensionContext, private readonly agentDir: string, private readonly modelId: string, private readonly thinking: ThinkingLevel, private readonly onEvent?: (event: MicroAgentEvent) => Promise<void>, private readonly onProviderFailure?: (failure: ProviderFailureClassification) => Promise<void>, private readonly cwd = context.cwd, private readonly allowedSkills: readonly string[] = []) {}

  async run(envelope: TaskEnvelope, options: { prompt?: string } = {}): Promise<AgentResult> {
    const [provider, modelId] = this.modelId.split("/", 2);
    const model = provider && modelId ? this.context.modelRegistry.find(provider, modelId) : undefined;
    if (!model) throw new Error(`Scout model is unavailable: ${this.modelId}`);
    const loader = await minimalResourceLoader(this.cwd, this.agentDir, envelope.scope, this.allowedSkills);
    const { session } = await createAgentSession({ cwd: this.cwd, agentDir: this.agentDir, model, thinkingLevel: this.thinking as never, tools: [...SCOPED_READ_TOOLS], customTools: scopedReadOnlyToolDefinitions(this.cwd, envelope.scope), resourceLoader: loader, sessionManager: SessionManager.inMemory(this.cwd) });
    let observedFailure: ProviderFailureClassification | undefined;
    let observedErrorClass: MicroAgentErrorClass | undefined;
    let terminalError = false;
    const observer = createMicroAgentEventObserver(this.modelId, this.thinking, this.onEvent);
    const unsubscribe = session.subscribe((event) => {
      observer.observe(event);
      if (event.type !== "message_end" || event.message.role !== "assistant") return;
      terminalError = event.message.stopReason === "error" || event.message.stopReason === "aborted";
      if (terminalError) {
        observedFailure = classifyProviderFailure(event.message.errorMessage);
        observedErrorClass = classifyMicroAgentFailure(event.message.errorMessage);
      }
    });
    const prompt = async (text: string): Promise<void> => {
      observedFailure = undefined;
      observedErrorClass = undefined;
      terminalError = false;
      const samplesBefore = observer.sampleCount();
      const statsBefore = session.getSessionStats();
      try {
        await session.prompt(text);
        await observer.flush();
      }
      catch (error) {
        const stats = session.getSessionStats();
        observer.failPrompt(error, { input: Math.max(0, stats.tokens.input - statsBefore.tokens.input), output: Math.max(0, stats.tokens.output - statsBefore.tokens.output), cacheRead: Math.max(0, stats.tokens.cacheRead - statsBefore.tokens.cacheRead), totalTokens: Math.max(0, stats.tokens.total - statsBefore.tokens.total), cost: Math.max(0, stats.cost - statsBefore.cost) }, samplesBefore);
        await observer.flush();
        const failure = classifyProviderFailure(error) ?? observedFailure;
        if (failure) await this.onProviderFailure?.(failure);
        throw error;
      }
      const failure = observedFailure as ProviderFailureClassification | undefined;
      if (failure) await this.onProviderFailure?.(failure);
      if (terminalError) throw new Error(`Worker provider failure: ${observedErrorClass ?? "micro-agent-failed"}`);
    };
    try {
      await prompt(options.prompt ?? envelopePrompt(envelope));
      try {
        return parseAgentResult(lastAssistantText(session.messages), envelope);
      } catch {
        await prompt("Your previous result violated the AgentResult contract. Submit exactly one corrected JSON AgentResult within the declared caps.");
        return parseAgentResult(lastAssistantText(session.messages), envelope);
      }
    } finally {
      unsubscribe();
      session.dispose();
    }
  }
}