import { ZodError } from "zod";
import { setTimeout as delay } from "node:timers/promises";
import { setDefaultAutoSelectFamilyAttemptTimeout } from "node:net";
import {
  acquireModel,
  cancelModelRequests,
  ModelError,
  providerCooldown,
  retryAfterMs,
} from "./model-runtime";
import type { Draft, Source, ExtractionMetrics } from "../src/types";
import { extractLocal } from "./context";
import { compactInput, parseCompactDraft, COMPACT_PROMPT } from "./compact";

type ModelConnection = {
  endpoint: string;
  key: string;
  model: string;
  auth: string;
};
// Node gives IPv6 250 ms before trying IPv4; api.meta.ai's IPv6 address can stall longer than that.
setDefaultAutoSelectFamilyAttemptTimeout(2_000);
let sessionConnection: ModelConnection | null = null;
let lastVerifiedAt: string | null = null;
let revision = 0;
const cache = new Map<string, { expires: number; draft: Draft }>();
const CACHE_TTL = 3 * 60 * 1000;
export const modelVerification = () => lastVerifiedAt;
export function configureAzureSession(connection: ModelConnection): void {
  const url = new URL(connection.endpoint);
  if (
    url.protocol !== "https:" ||
    url.username ||
    url.password ||
    url.port ||
    !/^[a-z0-9-]+\.(services\.ai|openai)\.azure\.com$/.test(url.hostname) ||
    url.pathname !== "/openai/v1/chat/completions" ||
    url.search ||
    url.hash
  ) {
    throw new Error(
      "Use the Azure deployment HTTPS chat-completions endpoint shown in Foundry.",
    );
  }
  if (
    !connection.key.trim() ||
    !/^[a-zA-Z0-9._-]{1,150}$/.test(connection.model) ||
    !["api-key", "bearer"].includes(connection.auth)
  ) {
    throw new Error(
      "The model connection fields are incomplete or unsupported.",
    );
  }
  cancelModelRequests();
  sessionConnection = { ...connection, key: connection.key.trim() };
  lastVerifiedAt = null;
  revision++;
  cache.clear();
}
export const modelConfig = () =>
  sessionConnection ?? {
    endpoint: process.env.CONTEXT_MODEL_ENDPOINT,
    key: process.env.CONTEXT_MODEL_API_KEY,
    model: process.env.CONTEXT_MODEL_ID,
    auth: process.env.CONTEXT_MODEL_AUTH || "bearer",
  };
export const modelReady = () => {
  const c = modelConfig();
  return Boolean(c.endpoint && c.key && c.model);
};
// Meta Model API (api.meta.ai) serves Muse Spark, a reasoning model: its thinking counts against the output cap.
export const modelProvider = (): "meta" | "azure" | "other" => {
  const endpoint = modelConfig().endpoint;
  if (!endpoint) return "other";
  const host = new URL(endpoint).hostname;
  return host === "api.meta.ai" ? "meta" : host.endsWith(".azure.com") ? "azure" : "other";
};
const MODEL_NAMES = { meta: "Meta Muse Spark", azure: "Azure model", other: "Model" };
export const modelLabel = () => (modelReady() ? `${MODEL_NAMES[modelProvider()]} (${modelConfig().model})` : null);

export async function extract(
  sources: Source[],
  useModel: boolean,
  freshModel = false,
  signal?: AbortSignal,
): Promise<Draft> {
  signal?.throwIfAborted();
  const started = performance.now(),
    draft = extractLocal(sources);
  const metrics: ExtractionMetrics = {
    providerCalls: 0,
    inputCharacters: 0,
    outputCharacters: 0,
    inputTokens: null,
    outputTokens: null,
    cacheHit: false,
  };
  const narrative = compactInput(sources);
  if (!useModel || narrative.length === 0)
    return { ...draft, strategy: "structured", attempts: 0, metrics };
  const c = modelConfig();
  if (!modelReady())
    throw new Error(
      "No model endpoint is configured. Choose the local structured reader.",
    );
  if (!/^https:\/\//.test(c.endpoint!))
    throw new Error("The model endpoint must use HTTPS.");
  const key = revision + ":" + c.model + ":" + draft.fingerprint;
  for (const [id, entry] of cache)
    if (entry.expires <= Date.now()) cache.delete(id);
  const hit = cache.get(key);
  if (!freshModel && hit)
    return {
      ...structuredClone(hit.draft),
      strategy: "cached",
      durationMs: Math.round(performance.now() - started),
      attempts: 0,
      metrics: { ...metrics, cacheHit: true },
    };
  const slot = acquireModel(signal);
  try {
    const input = JSON.stringify(narrative);
    metrics.inputCharacters = input.length;
    let usageReports = 0;
    const meta = modelProvider() === "meta";
    const deadline = performance.now() + (meta ? 60_000 : 22_000);
    let retryReason = "format";
    for (let attempt = 1; attempt <= 2; attempt++) {
      slot.signal.throwIfAborted();
      const remaining = Math.floor(deadline - performance.now());
      if (remaining < 500)
        throw new Error(
          "Extraction exceeded its time budget. No partial model response was used. Retry when the provider is available.",
        );
      let payload: {
        choices?: { finish_reason?: string; message?: { content?: unknown } }[];
        usage?: { prompt_tokens?: number; completion_tokens?: number };
      };
      try {
        metrics.providerCalls++;
        const response = await fetch(c.endpoint!, {
          method: "POST",
          signal: AbortSignal.any([
            slot.signal,
            AbortSignal.timeout(Math.min(meta ? 45_000 : 12_000, remaining)),
          ]),
          headers: {
            "Content-Type": "application/json",
            ...(c.auth === "api-key"
              ? { "api-key": c.key! }
              : { Authorization: "Bearer " + c.key }),
          },
          body: JSON.stringify({
            model: c.model,
            temperature: 0,
            ...(meta
              ? { max_completion_tokens: 12_000, reasoning_effort: "low" }
              : {
                  max_tokens: Math.min(
                    3000,
                    Math.max(500, draft.unparsed.length * 190),
                  ),
                }),
            response_format: { type: "json_object" },
            messages: [
              {
                role: "system",
                content:
                  COMPACT_PROMPT +
                  (attempt === 2 && retryReason === "format"
                    ? "\nReturn the exact JSON schema, with numeric references and one normalized string per fact. No prose."
                    : ""),
              },
              { role: "user", content: input },
            ],
          }),
        });
        if (
          [408, 429, 500, 502, 503, 504].includes(response.status) &&
          attempt === 1
        ) {
          const waitMs = retryAfterMs(response.headers.get("retry-after"));
          await response.body?.cancel();
          if (response.status === 429) providerCooldown(waitMs || 1000);
          if (waitMs > 2000)
            throw new ModelError(
              "The model provider is busy. Retry after the indicated delay; no model result was used.",
              503,
              Math.ceil(waitMs / 1000),
            );
          await delay(Math.max(150, waitMs), undefined, {
            signal: slot.signal,
          });
          retryReason = "transport";
          continue;
        }
        if (!response.ok) {
          const waitMs = retryAfterMs(response.headers.get("retry-after"));
          await response.body?.cancel();
          if (response.status === 429) providerCooldown(waitMs || 1000);
          throw new ModelError(
            response.status === 401 || response.status === 403
              ? "The provider rejected this connection. Check the model connection; no model result was used."
              : "The provider rejected the request or is busy. No model result was used. Retry later.",
            502,
            waitMs ? Math.ceil(waitMs / 1000) : undefined,
          );
        }
        payload = (await response.json()) as typeof payload;
      } catch (error) {
        if (slot.signal.aborted)
          throw new ModelError(
            "Extraction cancelled. No model result was used.",
            499,
          );
        if (error instanceof ModelError) throw error;
        if (attempt === 1) {
          retryReason = "transport";
          continue;
        }
        throw new Error(
          "The model request failed or timed out within the bounded retry budget. No model result was used.",
        );
      }
      slot.signal.throwIfAborted();
      if (
        !payload ||
        typeof payload !== "object" ||
        !Array.isArray(payload.choices)
      )
        throw new ModelError(
          "The provider returned an invalid response. No model result was used.",
          502,
        );
      const usage = payload.usage;
      if (
        usage &&
        Number.isSafeInteger(usage.prompt_tokens) &&
        Number.isSafeInteger(usage.completion_tokens) &&
        usage.prompt_tokens! >= 0 &&
        usage.completion_tokens! >= 0
      ) {
        usageReports++;
        metrics.inputTokens = (metrics.inputTokens ?? 0) + usage.prompt_tokens!;
        metrics.outputTokens =
          (metrics.outputTokens ?? 0) + usage.completion_tokens!;
      }
      const choice = payload.choices?.[0];
      if (choice?.finish_reason === "content_filter")
        throw new Error(
          "The provider's content filter blocked this input or response. No model result was used.",
        );
      if (choice?.finish_reason === "length")
        throw new Error(
          "The model reached its output limit. No model result was used. Split long records into smaller cases.",
        );
      const content = choice?.message?.content;
      if (typeof content !== "string" || !content.trim())
        throw new Error(
          "The provider returned no text. No model result was used.",
        );
      metrics.outputCharacters += content.length;
      let parsed: unknown;
      try {
        parsed = JSON.parse(
          content.trim().replace(/^```(?:json)?\s*([\s\S]*?)\s*```$/i, "$1"),
        );
      } catch {
        if (attempt === 1) {
          retryReason = "format";
          continue;
        }
        throw new Error(
          "The provider returned malformed JSON on both attempts. No model result was used.",
        );
      }
      let validated: Draft;
      try {
        validated = parseCompactDraft(
          parsed,
          sources,
          c.model!,
          Math.round(performance.now() - started),
        );
      } catch (error) {
        if (error instanceof ZodError && attempt === 1) {
          retryReason = "format";
          continue;
        }
        const reason =
          error instanceof ZodError
            ? "Unsupported model response structure."
            : error instanceof Error
              ? error.message
              : "Unsupported response.";
        throw new Error(
          "Model source checks failed. " +
            reason +
            " No model result was used.",
        );
      }
      lastVerifiedAt = new Date().toISOString();
      if (usageReports !== metrics.providerCalls) {
        metrics.inputTokens = null;
        metrics.outputTokens = null;
      }
      const result = { ...validated, attempts: attempt, metrics };
      if (cache.size >= 20) cache.delete(cache.keys().next().value!);
      cache.set(key, {
        expires: Date.now() + CACHE_TTL,
        draft: structuredClone(result),
      });
      return result;
    }
    throw new Error("No model result was used.");
  } catch (error) {
    if (slot.signal.aborted)
      throw new ModelError(
        "Extraction cancelled. No model result was used.",
        499,
      );
    if (error instanceof ModelError) throw error;
    throw new ModelError(
      error instanceof Error
        ? error.message
        : "Model extraction failed. No result was used.",
      502,
    );
  } finally {
    slot.release();
  }
}

const TRANSCRIBE_PROMPT =
  "Transcribe this synthetic medical document exactly as written, one output line per line on the page. " +
  "Keep the writer's words, numbers, units and abbreviations. Do not correct, interpret, summarise or add anything. " +
  "Write [illegible] for any word you cannot read. Return plain text only.";

/** Muse Spark reads a photo or PDF of a record into its lines; the reader then cites those lines like any typed record. */
export async function transcribe(data: Buffer, contentType: string, signal?: AbortSignal): Promise<string> {
  if (modelProvider() !== "meta") throw new ModelError("Photos and PDFs are read by Meta Muse Spark; connect it first.", 400);
  const c = modelConfig();
  const file =
    contentType === "application/pdf"
      ? { type: "file", file: { filename: "record.pdf", file_data: `data:application/pdf;base64,${data.toString("base64")}` } }
      : { type: "image_url", image_url: { url: `data:${contentType};base64,${data.toString("base64")}` } };
  const slot = acquireModel(signal);
  try {
    const response = await fetch(c.endpoint!, {
      method: "POST",
      signal: AbortSignal.any([slot.signal, AbortSignal.timeout(60_000)]),
      headers: { "Content-Type": "application/json", Authorization: "Bearer " + c.key },
      body: JSON.stringify({
        model: c.model,
        temperature: 0,
        max_completion_tokens: 8_000,
        reasoning_effort: "low",
        messages: [
          { role: "system", content: TRANSCRIBE_PROMPT },
          { role: "user", content: [{ type: "text", text: "Transcribe this record." }, file] },
        ],
      }),
    });
    if (!response.ok) {
      await response.body?.cancel();
      throw new ModelError("Muse Spark could not read this file. No reading was used.", 502);
    }
    const payload = (await response.json()) as { choices?: { message?: { content?: unknown } }[] };
    const text = payload.choices?.[0]?.message?.content;
    if (typeof text !== "string" || !text.trim()) throw new ModelError("Muse Spark returned no text for this file. No reading was used.", 502);
    return text.replace(/\r\n/g, "\n").replace(/^```\w*\n?|```$/g, "").trim();
  } finally {
    slot.release();
  }
}