import { env } from "cloudflare:workers";
import { backfillAccounts, backfillRepos } from "../backfill";
import { decryptSecret, encryptSecret } from "../crypto";
import { first } from "../db";
import { accountDb } from "../scoped-db";
import { badRequest, paymentRequired } from "../errors";
import { periodUsage } from "../metering/quota";
import { assertOutboundUrl } from "../integrations/outbound";

/**
 * Bring-your-own AI. An endpoint is saved per account or per repository; the repository's wins. Every AI
 * caller resolves its endpoint through resolveAiEndpoint, so precedence and validation live in one place.
 * Keys are sealed with SECRETS_KEY and never returned once saved (summaries carry no key material).
 */
export const AI_ENDPOINT_KINDS = ["provider", "openai", "cf-gateway"] as const;
export type AiEndpointKind = (typeof AI_ENDPOINT_KINDS)[number];
export type AiScope = "account" | "repo";

export const BYOK_PROVIDERS = [
  "openai",
  "anthropic",
  "google-ai-studio",
  "groq",
  "mistral",
  "deepseek",
  "cerebras",
  "grok",
  "openrouter",
  "workers-ai",
] as const;

export type AiEndpointInput = {
  kind: AiEndpointKind;
  /** Provider slug for 'provider' and 'cf-gateway' (routed through Gateway's OpenAI-compatible endpoint). */
  provider?: string | null;
  model: string;
  /** OpenAI-compatible base URL for 'openai', e.g. https://api.example.com/v1. */
  baseUrl?: string | null;
  cfAccountId?: string | null;
  cfGatewayId?: string | null;
  apiKey: string;
  /** Optional ceiling on the account's BYOK tokens per month (all its keys together); null for none. */
  monthlyTokenCap?: number | null;
};

export type AiEndpoint = Required<Omit<AiEndpointInput, "apiKey">> & {
  apiKey: string;
  scope: AiScope;
};

export type AiEndpointSummary = Omit<AiEndpoint, "apiKey"> & { updatedAt: number };

const MODEL = /^[\w.:@/-]{1,120}$/;
const CF_ACCOUNT = /^[a-f0-9]{32}$/;
const CF_GATEWAY = /^[a-z0-9][a-z0-9-]{0,63}$/;
const KEY_MIN = 8;
const KEY_MAX = 512;
const MAX_TOKEN_CAP = 1_000_000_000_000;

function tokenCap(value: number | null | undefined): number | null {
  if (value === null || value === undefined) return null;
  if (!Number.isSafeInteger(value) || value < 1 || value > MAX_TOKEN_CAP) throw badRequest("Invalid token cap");
  return value;
}

/** Normalizes and validates an endpoint; custom base URLs must be public https (loopback only in dev). */
export function validateEndpoint(input: AiEndpointInput): AiEndpointInput {
  const model = input.model.trim();
  const apiKey = input.apiKey.trim();
  if (!MODEL.test(model)) throw badRequest("Invalid model name");
  if (apiKey.length < KEY_MIN || apiKey.length > KEY_MAX) throw badRequest("Invalid API key");
  const provider = input.provider?.trim() || null;
  const monthlyTokenCap = tokenCap(input.monthlyTokenCap);
  if (input.kind !== "openai" && !BYOK_PROVIDERS.includes(provider as (typeof BYOK_PROVIDERS)[number])) {
    throw badRequest("Unknown provider");
  }
  if (input.kind === "openai") {
    const url = assertOutboundUrl((input.baseUrl ?? "").trim());
    if (url.search || url.hash) throw badRequest("The base URL may not have a query or fragment");
    const baseUrl = url.toString().replace(/\/+$/, "");
    return { kind: "openai", provider: null, model, apiKey, baseUrl, monthlyTokenCap };
  }
  if (input.kind === "cf-gateway") {
    const cfAccountId = (input.cfAccountId ?? "").trim().toLowerCase();
    const cfGatewayId = (input.cfGatewayId ?? "").trim().toLowerCase();
    if (!CF_ACCOUNT.test(cfAccountId)) throw badRequest("Cloudflare account id must be 32 hex characters");
    if (!CF_GATEWAY.test(cfGatewayId)) throw badRequest("Invalid AI Gateway id");
    return { kind: "cf-gateway", provider, model, apiKey, cfAccountId, cfGatewayId, monthlyTokenCap };
  }
  if (input.kind !== "provider") throw badRequest("Unknown endpoint kind");
  return { kind: "provider", provider, model, apiKey, monthlyTokenCap };
}

type EndpointRow = {
  scope: AiScope;
  kind: AiEndpointKind;
  provider: string | null;
  model: string;
  base_url: string | null;
  cf_account_id: string | null;
  cf_gateway_id: string | null;
  key_ciphertext: string;
  key_iv: string;
  scoped_seal: number;
  monthly_token_cap: number | null;
  updated_at: number;
};

function secretsKey(): string {
  if (!env.SECRETS_KEY) throw badRequest("SECRETS_KEY is not configured on this instance");
  return env.SECRETS_KEY;
}

function sealContext(scope: AiScope, ownerId: string) {
  return { purpose: "ai-endpoint", aad: `${scope}:${ownerId}` };
}

function summarize(row: EndpointRow): AiEndpointSummary {
  return {
    scope: row.scope,
    kind: row.kind,
    provider: row.provider,
    model: row.model,
    baseUrl: row.base_url,
    cfAccountId: row.cf_account_id,
    cfGatewayId: row.cf_gateway_id,
    monthlyTokenCap: row.monthly_token_cap,
    updatedAt: row.updated_at,
  };
}

/**
 * An endpoint row lives in the AccountState named by its owner id: the account's own object for account keys,
 * and an AccountState named by the repository id for repository keys (that is where the one-time import from
 * D1, which selects `owner_id = <object name>`, puts them).
 */
function endpointDb(ownerId: string) {
  return accountDb(ownerId);
}

async function loadRow(scope: AiScope, ownerId: string): Promise<EndpointRow | null> {
  return endpointDb(ownerId).first<EndpointRow>(
    "SELECT * FROM ai_endpoints WHERE scope = ? AND owner_id = ?",
    scope,
    ownerId,
  );
}

/** Legacy rows (generic seal from account_ai_keys) are re-sealed by resealLegacyEndpoints, never opened unbound. */
async function openRow(row: EndpointRow, ownerId: string): Promise<AiEndpoint> {
  if (row.scoped_seal === -1) throw badRequest("This AI key could not be upgraded; save it again");
  if (row.scoped_seal !== 1) throw badRequest("This AI key is being upgraded; try again in a few minutes");
  const apiKey = await decryptSecret(secretsKey(), row.key_ciphertext, row.key_iv, sealContext(row.scope, ownerId));
  const { updatedAt: _updatedAt, ...summary } = summarize(row);
  return { ...summary, apiKey };
}

const RESEAL_BATCH = 100;
const MAX_RESEAL_PASSES = 10;

async function resealOwner(ownerId: string): Promise<number> {
  const rows = await endpointDb(ownerId).all<EndpointRow & { owner_id: string }>(
    "SELECT * FROM ai_endpoints WHERE owner_id = ? AND scoped_seal = 0 ORDER BY updated_at LIMIT ?",
    ownerId,
    RESEAL_BATCH,
  );
  for (const row of rows) {
    const apiKey = await decryptSecret(secretsKey(), row.key_ciphertext, row.key_iv).catch((error: unknown) => {
      console.error("Couldn't open a legacy AI key; it must be re-saved", row.scope, row.owner_id, error);
      return null;
    });
    if (apiKey === null) {
      await endpointDb(row.owner_id).run(
        "UPDATE ai_endpoints SET scoped_seal = -1 WHERE scope = ? AND owner_id = ? AND scoped_seal = 0",
        row.scope,
        row.owner_id,
      );
      continue;
    }
    const sealed = await encryptSecret(secretsKey(), apiKey, sealContext(row.scope, row.owner_id));
    await endpointDb(row.owner_id).run(
      `UPDATE ai_endpoints SET key_ciphertext = ?, key_iv = ?, scoped_seal = 1
       WHERE scope = ? AND owner_id = ? AND key_ciphertext = ? AND scoped_seal = 0`,
      sealed.ciphertext,
      sealed.iv,
      row.scope,
      row.owner_id,
      row.key_ciphertext,
    );
  }
  return rows.length;
}

/**
 * Cron until done: re-seals rows copied from account_ai_keys to their scope and owner (AAD), so no unbound seal
 * stays. Walks accounts, then repositories (whose endpoints live in an AccountState named by the repository id),
 * a page per run; new keys are always sealed to their scope, so once both walks finish nothing is left to do.
 */
export async function resealLegacyEndpoints(): Promise<void> {
  const reseal = async (ownerId: string) => {
    // Rows whose conditional update lost a race stay unsealed; a few passes is plenty, never loop forever.
    for (let pass = 0; pass < MAX_RESEAL_PASSES; pass++) {
      if ((await resealOwner(ownerId)) < RESEAL_BATCH) return;
    }
  };
  if (await backfillAccounts("ai-reseal-accounts", reseal)) await backfillRepos("ai-reseal-repos", reseal);
}

export async function endpointSummary(scope: AiScope, ownerId: string): Promise<AiEndpointSummary | null> {
  const row = await loadRow(scope, ownerId);
  return row ? summarize(row) : null;
}

export async function saveEndpoint(
  scope: AiScope,
  ownerId: string,
  input: AiEndpointInput,
  actorId: string,
): Promise<AiEndpointSummary> {
  const valid = validateEndpoint(input);
  const sealed = await encryptSecret(secretsKey(), valid.apiKey, sealContext(scope, ownerId));
  const now = Date.now();
  await endpointDb(ownerId).run(
    `INSERT INTO ai_endpoints (scope, owner_id, kind, provider, model, base_url, cf_account_id, cf_gateway_id,
       key_ciphertext, key_iv, updated_by, updated_at, monthly_token_cap) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
     ON CONFLICT (scope, owner_id) DO UPDATE SET kind = excluded.kind, provider = excluded.provider,
       model = excluded.model, base_url = excluded.base_url, cf_account_id = excluded.cf_account_id,
       cf_gateway_id = excluded.cf_gateway_id, key_ciphertext = excluded.key_ciphertext, key_iv = excluded.key_iv, scoped_seal = 1,
       updated_by = excluded.updated_by, updated_at = excluded.updated_at, monthly_token_cap = excluded.monthly_token_cap`,
    scope,
    ownerId,
    valid.kind,
    valid.provider ?? null,
    valid.model,
    valid.baseUrl ?? null,
    valid.cfAccountId ?? null,
    valid.cfGatewayId ?? null,
    sealed.ciphertext,
    sealed.iv,
    actorId,
    now,
    valid.monthlyTokenCap ?? null,
  );
  return {
    scope,
    kind: valid.kind,
    provider: valid.provider ?? null,
    model: valid.model,
    baseUrl: valid.baseUrl ?? null,
    cfAccountId: valid.cfAccountId ?? null,
    cfGatewayId: valid.cfGatewayId ?? null,
    monthlyTokenCap: valid.monthlyTokenCap ?? null,
    updatedAt: now,
  };
}

export async function deleteEndpoint(scope: AiScope, ownerId: string): Promise<void> {
  await endpointDb(ownerId).run("DELETE FROM ai_endpoints WHERE scope = ? AND owner_id = ?", scope, ownerId);
}

/**
 * A repository's endpoint, but only when the billed account owns that repository: callers that bill the viewer
 * (pull descriptions) must never spend a repository owner's key.
 */
async function ownedRepoRow(repoId: string, accountId: string): Promise<EndpointRow | null> {
  const owned = await first<{ ok: number }>(
    "SELECT 1 AS ok FROM repos WHERE id = ? AND owner_id = ?",
    repoId,
    accountId,
  );
  if (!owned) return null;
  return endpointDb(repoId).first<EndpointRow>(
    "SELECT * FROM ai_endpoints WHERE scope = 'repo' AND owner_id = ?",
    repoId,
  );
}

/** A key with a monthly cap stops being used once the account's BYOK tokens this month reach it. */
async function withinCap(row: EndpointRow, accountId: string): Promise<EndpointRow> {
  if (row.monthly_token_cap === null) return row;
  if ((await periodUsage(accountId, "ai.byok_tokens")).quantity < row.monthly_token_cap) return row;
  throw paymentRequired("This AI key's monthly token cap is reached", {
    meter: "ai.byok_tokens",
    reason: "spending_limit",
    accountId,
  });
}

/**
 * The endpoint an AI call billed to `accountId` must use (the repository's, else the account's), or null for
 * the platform model. The single resolver for every AI caller.
 */
export async function resolveAiEndpoint(context: {
  accountId: string;
  repoId?: string | null;
}): Promise<AiEndpoint | null> {
  if (context.repoId) {
    const repo = await ownedRepoRow(context.repoId, context.accountId);
    if (repo) return openRow(await withinCap(repo, context.accountId), context.repoId);
  }
  const account = await loadRow("account", context.accountId);
  return account ? openRow(await withinCap(account, context.accountId), context.accountId) : null;
}

/** Same precedence as resolveAiEndpoint, without decrypting: for access checks and settings pages. */
export async function resolveAiEndpointSummary(context: {
  accountId: string;
  repoId?: string | null;
}): Promise<AiEndpointSummary | null> {
  if (context.repoId) {
    const repo = await ownedRepoRow(context.repoId, context.accountId);
    if (repo) return summarize(repo);
  }
  return endpointSummary("account", context.accountId);
}