import { Effect, Schema } from "effect";
import { fetchInference } from "../../http/local-fetch";
import type { AppContext } from "../../app-context";
import type { Recipe } from "../models/types";
import type { ProxyDialect } from "./usage-observer";
import type { InferenceUsageInput } from "./inference-accounting";

/**
 * Context guard: reject local inference requests once a session's known
 * context depth crosses a soft ceiling below the engine's max_model_len.
 *
 * Local engines degrade — and experimental builds can die outright — well
 * before the hard context limit, and a request that sails past the fragile
 * zone takes the whole engine down for every client. Agent clients (pi,
 * ZCode, and most OpenAI-compatible SDKs) already recognise context-overflow
 * errors and respond by compacting the conversation and retrying, so the
 * kindest thing the proxy can do is fail the request in exactly that error
 * shape while the engine is still healthy. The rejection message deliberately
 * contains "maximum context length is N tokens" and "context_length_exceeded"
 * (OpenAI dialects) or "prompt is too long" (Anthropic dialect) — the phrases
 * overflow detectors match on.
 *
 * Depth is judged from the engine's own numbers: after every completion the
 * proxy records prompt+completion tokens per session id, and the next request
 * from that session is checked against the ceiling. Requests without a
 * session id fall back to a conservative size estimate of the request body.
 */

const WINDOW_CACHE_TTL_MS = 60_000;
const SESSION_TTL_MS = 6 * 60 * 60 * 1000;
const MAX_TRACKED_SESSIONS = 1000;
/** Rough chars-per-token for the estimate fallback; high (= conservative)
 * so serialization overhead and code-dense text don't cause false rejects. */
const ESTIMATE_CHARS_PER_TOKEN = 4;
/** Estimates only reject with extra margin — they are a backstop for
 * sessions the proxy has no exact usage record for. */
const ESTIMATE_SLACK = 1.15;

const ActiveModelsSchema = Schema.Struct({
  data: Schema.optional(
    Schema.Array(Schema.Struct({ max_model_len: Schema.optional(Schema.Number) })),
  ),
});

export interface ContextLimits {
  window: number;
  softLimit: number;
}

export interface ContextRejection {
  status: 400;
  body: Record<string, unknown>;
}

interface UsageRecordKey {
  model: string;
  session_id: string | null;
  provider: string;
}

const sessionKey = (model: string, sessionId: string): string => `${model} ${sessionId}`;

const formatTokens = (n: number): string => String(Math.round(n));

const rejectionBody = (
  dialect: ProxyDialect,
  contextTokens: number,
  limits: ContextLimits,
): Record<string, unknown> => {
  if (dialect === "messages") {
    const message =
      `prompt is too long: ~${formatTokens(contextTokens)} tokens > ` +
      `${formatTokens(limits.softLimit)} soft maximum (model maximum ${formatTokens(limits.window)}). ` +
      `Compact or summarize the conversation and retry.`;
    return { type: "error", error: { type: "invalid_request_error", message } };
  }
  const message =
    `Session context is ~${formatTokens(contextTokens)} tokens; the soft limit for this model is ` +
    `${formatTokens(limits.softLimit)} tokens and its maximum context length is ` +
    `${formatTokens(limits.window)} tokens. Compact or summarize the conversation and retry. ` +
    `(context_length_exceeded)`;
  return {
    error: {
      message,
      type: "invalid_request_error",
      param: "messages",
      code: "context_length_exceeded",
    },
  };
};

/** Serialized size of the conversation content, ignoring inline images and
 * other base64 payloads (they are huge in chars but small in tokens). */
const estimateBodyTokens = (dialect: ProxyDialect, body: Record<string, unknown>): number => {
  const content =
    dialect === "responses"
      ? [body["input"], body["instructions"]]
      : [body["messages"], body["system"]];
  let chars = 0;
  for (const part of content) {
    if (part === undefined || part === null) continue;
    try {
      chars += JSON.stringify(part, (_key, value: unknown): unknown =>
        typeof value === "string" && looksLikeBase64(value) ? "" : value,
      ).length;
    } catch {
      // Circular or otherwise unserializable content — skip the estimate.
      return 0;
    }
  }
  return Math.ceil(chars / ESTIMATE_CHARS_PER_TOKEN);
};

/** Inline binary payloads only: an explicit data: URL, or a very long
 * unbroken base64-alphabet run. Prose and code both break the run with
 * spaces or punctuation long before the length threshold. */
const looksLikeBase64 = (value: string): boolean => {
  if (value.startsWith("data:")) return /^data:[\w/+.-]+;base64,/.test(value.slice(0, 128));
  return value.length > 16384 && /^[A-Za-z0-9+/=_-]{16384}/.test(value.slice(0, 16384));
};

export interface ContextGuard {
  resolveLimits: (recipe: Recipe) => Effect.Effect<ContextLimits | null>;
  observe: (record: UsageRecordKey, usage: InferenceUsageInput) => void;
  check: (input: {
    dialect: ProxyDialect;
    sessionId: string | null;
    model: string;
    body: Record<string, unknown>;
    limits: ContextLimits;
  }) => ContextRejection | null;
  headers: (limits: ContextLimits) => Record<string, string>;
}

export const createContextGuard = (context: AppContext): ContextGuard => {
  const sessions = new Map<string, { tokens: number; at: number }>();
  let windowCache: { recipeId: string; value: number | null; at: number } | null = null;

  const pruneSessions = (): void => {
    if (sessions.size <= MAX_TRACKED_SESSIONS) return;
    const cutoff = Date.now() - SESSION_TTL_MS;
    for (const [key, entry] of sessions) {
      if (entry.at < cutoff) sessions.delete(key);
    }
    // Still oversized after TTL pruning: drop oldest insertions (Map iterates
    // in insertion order, and re-inserts on update keep hot sessions late).
    while (sessions.size > MAX_TRACKED_SESSIONS) {
      const oldest = sessions.keys().next();
      if (oldest.done) break;
      sessions.delete(oldest.value);
    }
  };

  const fetchLiveWindow = (): Effect.Effect<number | null> =>
    fetchInference(context, "/v1/models", { timeoutMs: 3000 }).pipe(
      Effect.flatMap((response) =>
        Effect.tryPromise({ try: () => response.json(), catch: (error) => error }),
      ),
      Effect.flatMap(Schema.decodeUnknownEffect(ActiveModelsSchema)),
      Effect.map((decoded): number | null => decoded.data?.[0]?.max_model_len ?? null),
      Effect.catch(() => Effect.succeed(null)),
    );

  /** max_model_len for the recipe: the recipe's own value, else the live
   * engine's /v1/models report, cached briefly. Null disables the guard. */
  const resolveLimits = (recipe: Recipe): Effect.Effect<ContextLimits | null> =>
    Effect.gen(function* () {
      const ratio = context.config.context_soft_limit_ratio;
      if (!(ratio > 0) || ratio >= 1) return null;
      let window: number | null = recipe.max_model_len ?? null;
      if (!window) {
        if (windowCache?.recipeId === recipe.id && Date.now() - windowCache.at < WINDOW_CACHE_TTL_MS) {
          window = windowCache.value;
        } else {
          window = yield* fetchLiveWindow();
          windowCache = { recipeId: recipe.id, value: window, at: Date.now() };
        }
      }
      if (!window || window <= 0) return null;
      return { window, softLimit: Math.floor(window * ratio) };
    });

  /** Called with the engine-reported usage of every completed local request. */
  const observe = (record: UsageRecordKey, usage: InferenceUsageInput): void => {
    if (record.provider !== "local" || !record.session_id) return;
    const prompt = usage.prompt_tokens ?? 0;
    const completion = usage.completion_tokens ?? 0;
    const tokens = prompt + completion;
    if (tokens <= 0) return;
    const key = sessionKey(record.model, record.session_id);
    sessions.delete(key);
    sessions.set(key, { tokens, at: Date.now() });
    pruneSessions();
  };

  /** Null when the request may pass; a dialect-shaped 400 when it must compact. */
  const check = (input: {
    dialect: ProxyDialect;
    sessionId: string | null;
    model: string;
    body: Record<string, unknown>;
    limits: ContextLimits;
  }): ContextRejection | null => {
    const { limits } = input;
    if (input.sessionId) {
      const known = sessions.get(sessionKey(input.model, input.sessionId));
      if (known && Date.now() - known.at < SESSION_TTL_MS && known.tokens > limits.softLimit) {
        return { status: 400, body: rejectionBody(input.dialect, known.tokens, limits) };
      }
    }
    const estimate = estimateBodyTokens(input.dialect, input.body);
    if (estimate > limits.softLimit * ESTIMATE_SLACK) {
      return { status: 400, body: rejectionBody(input.dialect, estimate, limits) };
    }
    return null;
  };

  const headers = (limits: ContextLimits): Record<string, string> => ({
    "x-context-window": String(limits.window),
    "x-context-soft-limit": String(limits.softLimit),
  });

  return { resolveLimits, observe, check, headers };
};