import { completeSimple, getSupportedThinkingLevels } from "@earendil-works/pi-ai/compat";
import type {
  AssistantMessage,
  Context,
  Model,
  ModelThinkingLevel,
  SimpleStreamOptions,
} from "@earendil-works/pi-ai";
import type { ModelRegistry as PiModelRegistry } from "@earendil-works/pi-coding-agent";
import type { Assessment } from "./assessment.js";
import type { ReviewerConfig } from "./config.js";
import { parseAssessment } from "./assessment.js";
import { GUARDIAN_POLICY } from "./prompt.js";

export const REVIEW_DEADLINE_MS = 90_000;
export const REVIEW_MAX_ATTEMPTS = 3;
export const REVIEW_MAX_TOKENS = 384;

export type ReviewResult =
  | { status: "allowed"; assessment: Assessment; attempts: number }
  | { status: "denied"; assessment: Assessment; attempts: number }
  | { status: "failed-closed"; reason: string; attempts: number }
  | { status: "timed-out"; reason: string; attempts: number }
  | { status: "aborted"; reason: string; attempts: number };

/** Narrow public Pi registry contract used by the reviewer and deterministic tests. */
export type ModelRegistry = Pick<
  PiModelRegistry,
  "find" | "getApiKeyAndHeaders" | "getRegisteredProviderConfig"
>;

export type CompleteSimple = (
  model: Model<any>,
  context: Context,
  options: SimpleStreamOptions,
) => Promise<AssistantMessage>;

export interface ReviewerDependencies {
  complete: CompleteSimple;
  getSupportedReasoning: (model: Model<any>) => ModelThinkingLevel[];
  now: () => number;
  setTimer: (callback: () => void, milliseconds: number) => ReturnType<typeof setTimeout>;
  clearTimer: (timer: ReturnType<typeof setTimeout>) => void;
}

const DEFAULT_DEPENDENCIES: ReviewerDependencies = {
  complete: completeSimple,
  getSupportedReasoning: getSupportedThinkingLevels,
  now: Date.now,
  setTimer: setTimeout,
  clearTimer: clearTimeout,
};

class ReviewAbort extends Error {}

function abortable<T>(promise: Promise<T>, signal: AbortSignal): Promise<T> {
  if (signal.aborted) return Promise.reject(new ReviewAbort("Review aborted"));
  return new Promise<T>((resolve, reject) => {
    const onAbort = (): void => reject(new ReviewAbort("Review aborted"));
    signal.addEventListener("abort", onAbort, { once: true });
    promise.then(
      (value) => {
        signal.removeEventListener("abort", onAbort);
        resolve(value);
      },
      (error: unknown) => {
        signal.removeEventListener("abort", onAbort);
        reject(error);
      },
    );
  });
}

function sleep(milliseconds: number, signal: AbortSignal, deps: ReviewerDependencies): Promise<void> {
  return new Promise<void>((resolve, reject) => {
    if (signal.aborted) {
      reject(new ReviewAbort("Review aborted"));
      return;
    }
    const timer = deps.setTimer(() => {
      signal.removeEventListener("abort", onAbort);
      resolve();
    }, milliseconds);
    const onAbort = (): void => {
      deps.clearTimer(timer);
      reject(new ReviewAbort("Review aborted"));
    };
    signal.addEventListener("abort", onAbort, { once: true });
  });
}

function textOutput(message: AssistantMessage): string {
  return message.content
    .filter((item) => item.type === "text")
    .map((item) => item.text)
    .join("\n");
}

function terminalAbortResult(
  cause: "parent" | "deadline" | undefined,
  attempts: number,
): Extract<ReviewResult, { status: "timed-out" | "aborted" | "failed-closed" }> {
  if (cause === "deadline") {
    return { status: "timed-out", reason: "Automatic approval review exceeded its 90-second deadline.", attempts };
  }
  if (cause === "parent") {
    return { status: "aborted", reason: "Automatic approval review was cancelled with the active Pi turn.", attempts };
  }
  return { status: "failed-closed", reason: "The reviewer stopped without a final assessment.", attempts };
}

export async function reviewWithModel(
  config: ReviewerConfig,
  prompt: string,
  registry: ModelRegistry,
  parentSignal?: AbortSignal,
  overrides: Partial<ReviewerDependencies> = {},
  deadlineMs = REVIEW_DEADLINE_MS,
): Promise<ReviewResult> {
  const deps = { ...DEFAULT_DEPENDENCIES, ...overrides };
  const controller = new AbortController();
  let abortCause: "parent" | "deadline" | undefined;
  let attempts = 0;
  const deadlineAt = deps.now() + deadlineMs;

  const onParentAbort = (): void => {
    if (abortCause === undefined) abortCause = "parent";
    controller.abort();
  };
  if (parentSignal?.aborted === true) onParentAbort();
  else parentSignal?.addEventListener("abort", onParentAbort, { once: true });

  const deadlineTimer = deps.setTimer(() => {
    if (abortCause === undefined) abortCause = "deadline";
    controller.abort();
  }, deadlineMs);

  try {
    for (let attempt = 1; attempt <= REVIEW_MAX_ATTEMPTS; attempt += 1) {
      attempts = attempt;
      if (controller.signal.aborted || deps.now() >= deadlineAt) {
        if (abortCause === undefined) abortCause = "deadline";
        controller.abort();
        return terminalAbortResult(abortCause, attempts);
      }

      const slash = config.model.indexOf("/");
      const provider = config.model.slice(0, slash);
      const modelId = config.model.slice(slash + 1);
      const model = registry.find(provider, modelId);
      if (model === undefined) {
        return { status: "failed-closed", reason: `Configured reviewer model ${config.model} is not in the current model registry.`, attempts };
      }
      if (!deps.getSupportedReasoning(model).includes(config.reasoning)) {
        return {
          status: "failed-closed",
          reason: `Reviewer model ${config.model} does not support configured reasoning ${config.reasoning}.`,
          attempts,
        };
      }

      let auth: Awaited<ReturnType<ModelRegistry["getApiKeyAndHeaders"]>>;
      try {
        auth = await abortable(registry.getApiKeyAndHeaders(model), controller.signal);
      } catch {
        if (controller.signal.aborted) return terminalAbortResult(abortCause, attempts);
        if (attempt < REVIEW_MAX_ATTEMPTS) {
          try {
            await sleep(200 * 2 ** (attempt - 1), controller.signal, deps);
            continue;
          } catch {
            return terminalAbortResult(abortCause, attempts);
          }
        }
        return { status: "failed-closed", reason: "Reviewer authentication resolution failed.", attempts };
      }
      if (!auth.ok) {
        return { status: "failed-closed", reason: `Reviewer authentication is unavailable: ${auth.error}`, attempts };
      }

      const remaining = Math.max(1, deadlineAt - deps.now());
      const options: SimpleStreamOptions = {
        reasoning: config.reasoning,
        signal: controller.signal,
        maxTokens: REVIEW_MAX_TOKENS,
        maxRetries: 0,
        timeoutMs: remaining,
      };
      if (auth.apiKey !== undefined) options.apiKey = auth.apiKey;
      if (auth.headers !== undefined) options.headers = auth.headers;
      if (auth.env !== undefined) options.env = auth.env;

      const context: Context = {
        systemPrompt: GUARDIAN_POLICY,
        messages: [{ role: "user", content: prompt, timestamp: deps.now() }],
        tools: [],
      };
      const registeredProvider = registry.getRegisteredProviderConfig(provider);
      const registeredStreamSimple =
        registeredProvider?.api === model.api ? registeredProvider.streamSimple : undefined;

      let response: AssistantMessage;
      try {
        const completion = registeredStreamSimple === undefined
          ? deps.complete(model, context, options)
          : registeredStreamSimple(model, context, options).result();
        response = await abortable(completion, controller.signal);
      } catch {
        if (controller.signal.aborted) return terminalAbortResult(abortCause, attempts);
        if (attempt < REVIEW_MAX_ATTEMPTS) {
          try {
            await sleep(200 * 2 ** (attempt - 1), controller.signal, deps);
            continue;
          } catch {
            return terminalAbortResult(abortCause, attempts);
          }
        }
        return { status: "failed-closed", reason: "Reviewer provider requests failed after three attempts.", attempts };
      }

      if (response.stopReason === "aborted") return terminalAbortResult(abortCause, attempts);
      if (response.stopReason === "error") {
        if (attempt < REVIEW_MAX_ATTEMPTS) {
          try {
            await sleep(200 * 2 ** (attempt - 1), controller.signal, deps);
            continue;
          } catch {
            return terminalAbortResult(abortCause, attempts);
          }
        }
        return { status: "failed-closed", reason: "Reviewer provider requests failed after three attempts.", attempts };
      }
      if (response.stopReason !== "stop") {
        return {
          status: "failed-closed",
          reason: `Reviewer returned non-final stop reason ${response.stopReason}.`,
          attempts,
        };
      }

      try {
        const assessment = parseAssessment(textOutput(response));
        return assessment.outcome === "allow"
          ? { status: "allowed", assessment, attempts }
          : { status: "denied", assessment, attempts };
      } catch {
        if (attempt < REVIEW_MAX_ATTEMPTS) {
          try {
            await sleep(200 * 2 ** (attempt - 1), controller.signal, deps);
          } catch {
            return terminalAbortResult(abortCause, attempts);
          }
          continue;
        }
        return { status: "failed-closed", reason: "Reviewer output was malformed after three attempts.", attempts };
      }
    }
    throw new Error("Unreachable reviewer attempt state");
  } finally {
    deps.clearTimer(deadlineTimer);
    parentSignal?.removeEventListener("abort", onParentAbort);
  }
}