import { env, waitUntil } from "cloudflare:workers";
import { HttpError } from "../errors";
import { limitMcpRequest } from "./audit";
import { recordToolUsage } from "./usage";
import { authenticate, agentIdentity, unauthorizedResponse } from "./caller";
import type { Authenticated } from "./caller";
import { dispatch, finishResult, initializeResult, negotiateLegacyVersion } from "./dispatch";
import { sessionHub } from "./hub-do";
import type { SessionInfo } from "./hub-session";
import { sseFrame } from "./hub-session";
import { openListenStream } from "./listen";
import {
  isObject,
  META,
  toJsonText,
  MODERN_VERSION,
  notification,
  RPC,
  rpcError,
  RpcError,
  rpcResult,
  SUPPORTED_VERSIONS,
} from "./protocol";
import type { JsonObject, JsonRpcId, JsonRpcMessage, McpCaller } from "./protocol";

const MAX_BODY_BYTES = 4 * 1024 * 1024;
const HTTP_OK = 200;
const HTTP_BAD_REQUEST = 400;
const HTTP_TOO_LARGE = 413;
const SESSION_HEADER = "mcp-session-id";
const NAMED_METHODS: Record<string, string> = { "tools/call": "name", "prompts/get": "name", "resources/read": "uri" };
const encoder = new TextEncoder();

function jsonResponse(message: JsonRpcMessage, status = 200, headers: Record<string, string> = {}): Response {
  return Response.json(message, { status, headers: { "cache-control": "no-store", ...headers } });
}

function errorResponse(
  id: JsonRpcId | null,
  error: RpcError,
  status = error.httpStatus === HTTP_OK ? HTTP_BAD_REQUEST : error.httpStatus,
) {
  return jsonResponse(rpcError(id, error), status);
}

/** Reads at most `limit` bytes of the body; null when it is larger, whatever Content-Length claims. */
async function readBounded(request: Request, limit: number): Promise<string | null> {
  if (Number(request.headers.get("content-length") ?? "0") > limit) return null;
  if (!request.body) return "";
  const reader = request.body.getReader();
  const chunks: Uint8Array[] = [];
  let size = 0;
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    size += value.byteLength;
    if (size > limit) {
      await reader.cancel().catch(() => undefined);
      return null;
    }
    chunks.push(value);
  }
  const bytes = new Uint8Array(size);
  let offset = 0;
  for (const chunk of chunks) {
    bytes.set(chunk, offset);
    offset += chunk.byteLength;
  }
  return new TextDecoder().decode(bytes);
}

function decodeHeaderValue(raw: string): string {
  const match = /^=\?base64\?(.*)\?=$/.exec(raw);
  if (!match) return raw;
  try {
    return new TextDecoder().decode(Uint8Array.from(atob(match[1]), (char) => char.charCodeAt(0)));
  } catch {
    return "\u0000invalid";
  }
}

/** 2026-07-28 requires routing headers that mirror the body; any disagreement is a HeaderMismatch. */
function checkModernHeaders(request: Request, method: string, params: JsonObject, version: string): RpcError | null {
  const mismatch = (message: string) => new RpcError(RPC.headerMismatch, `Header mismatch: ${message}`);
  if (request.headers.get("mcp-protocol-version") !== version) return mismatch("MCP-Protocol-Version");
  if (request.headers.get("mcp-method") !== method) return mismatch("Mcp-Method");
  const nameKey = NAMED_METHODS[method];
  if (nameKey) {
    const header = request.headers.get("mcp-name");
    if (header === null || decodeHeaderValue(header) !== params[nameKey]) return mismatch("Mcp-Name");
  }
  return null;
}

function metaOf(params: JsonObject): JsonObject {
  return isObject(params._meta) ? params._meta : {};
}

async function buildCaller(
  auth: Authenticated,
  request: Request,
  url: URL,
  era: McpCaller["era"],
  session: { id: string; info: SessionInfo } | null,
  params: JsonObject,
): Promise<McpCaller> {
  const meta = metaOf(params);
  const clientInfo = isObject(meta[META.clientInfo]) ? meta[META.clientInfo] : null;
  const declared = session ? JSON.parse(session.info.clientCapabilities) : meta[META.clientCapabilities];
  const capabilities = isObject(declared) ? declared : {};
  return {
    credential: auth.credential,
    viewer: auth.viewer,
    tokenId: auth.tokenId,
    agent: session?.info.agent ?? (await agentIdentity(auth, request, isObject(clientInfo) ? clientInfo.name : null)),
    era,
    sessionId: session?.id ?? null,
    clientCapabilities: capabilities,
    origin: url.origin,
  };
}

async function loadSession(auth: Authenticated, sessionId: string): Promise<SessionInfo | null> {
  if (!/^[A-Za-z0-9-]{16,64}$/.test(sessionId)) return null;
  const info = await sessionHub(sessionId).session();
  if (!info || info.kind !== "legacy" || info.userId !== auth.viewer.id || info.tokenId !== auth.tokenId) return null;
  return info;
}

const sessionMissing = () =>
  errorResponse(null, new RpcError(RPC.invalidRequest, "Session not found; initialize again"), 404);

async function initialize(auth: Authenticated, request: Request, id: JsonRpcId, params: JsonObject) {
  // Each initialize creates a Durable Object, storage and an alarm: budget it like a write (INT-A12).
  try {
    await limitMcpRequest(auth.tokenId, null, true);
  } catch (error) {
    if (error instanceof HttpError)
      return errorResponse(id, new RpcError(RPC.invalidRequest, error.message), error.status);
    throw error;
  }
  waitUntil(recordToolUsage(auth.viewer.id, "$initialize", false));
  const protocolVersion = negotiateLegacyVersion(params.protocolVersion);
  const clientInfo = isObject(params.clientInfo) ? params.clientInfo : {};
  const sessionId = crypto.randomUUID();
  const info: SessionInfo = {
    kind: "legacy",
    userId: auth.viewer.id,
    tokenId: auth.tokenId,
    agent: await agentIdentity(auth, request, clientInfo.name),
    protocolVersion,
    clientCapabilities: toJsonText(isObject(params.capabilities) ? params.capabilities : {}),
    createdAt: Date.now(),
    subscriptionId: null,
    resourceUris: [],
    repos: [],
    toolEvents: {},
  };
  await sessionHub(sessionId).openSession(info);
  return jsonResponse(rpcResult(id, initializeResult(protocolVersion)), 200, { [SESSION_HEADER]: sessionId });
}

/** Answers on an SSE stream when the client asked for progress, so long tools can report as they go. */
function streamed(caller: McpCaller, id: JsonRpcId, method: string, params: JsonObject, request: Request): Response {
  const progressToken = metaOf(params).progressToken;
  const { readable, writable } = new TransformStream<Uint8Array, Uint8Array>();
  const writer = writable.getWriter();
  const send = (message: JsonRpcMessage) =>
    writer.write(encoder.encode(sseFrame(JSON.stringify(message)))).catch(() => undefined);
  const progress = (message: string, value: number, total?: number) => {
    void send(
      notification("notifications/progress", { progressToken, progress: value, ...(total ? { total } : {}), message }),
    );
  };
  const work = dispatch(caller, method, params, { signal: request.signal, progress })
    .then((result) => send(rpcResult(id, finishResult(caller, result))))
    .catch((error: unknown) =>
      send(rpcError(id, error instanceof RpcError ? error : new RpcError(RPC.internal, "Internal error"))),
    )
    .finally(() => writer.close().catch(() => undefined));
  void work;
  return new Response(readable, {
    headers: {
      "content-type": "text/event-stream",
      "cache-control": "no-cache, no-transform",
      "x-accel-buffering": "no",
    },
  });
}

async function respond(caller: McpCaller, request: Request, id: JsonRpcId, method: string, params: JsonObject) {
  const wantsStream = (request.headers.get("accept") ?? "").includes("text/event-stream");
  if (wantsStream && method === "tools/call" && metaOf(params).progressToken !== undefined) {
    return streamed(caller, id, method, params, request);
  }
  try {
    const result = await dispatch(caller, method, params, { signal: request.signal, progress: () => undefined });
    return jsonResponse(rpcResult(id, finishResult(caller, result)));
  } catch (error) {
    if (error instanceof RpcError) return errorResponse(id, error, error.httpStatus);
    console.error("MCP request failed", method, error);
    return errorResponse(id, new RpcError(RPC.internal, "Internal error"), 500);
  }
}

async function handlePost(auth: Authenticated, request: Request, url: URL): Promise<Response> {
  const text = await readBounded(request, MAX_BODY_BYTES);
  if (text === null) return errorResponse(null, new RpcError(RPC.invalidRequest, "Body too large"), HTTP_TOO_LARGE);
  let message: unknown;
  try {
    message = JSON.parse(text);
  } catch {
    return errorResponse(null, new RpcError(RPC.parse, "Parse error"));
  }
  if (!isObject(message) || message.jsonrpc !== "2.0") {
    return errorResponse(null, new RpcError(RPC.invalidRequest, "Expected a single JSON-RPC 2.0 message"));
  }
  const sessionId = request.headers.get(SESSION_HEADER);
  const method = typeof message.method === "string" ? message.method : null;
  const id = typeof message.id === "string" || typeof message.id === "number" ? message.id : null;
  const params = isObject(message.params) ? message.params : {};
  if (!method || id === null) {
    if (sessionId && !(await loadSession(auth, sessionId))) return sessionMissing();
    return new Response(null, { status: 202 });
  }
  if (method === "initialize") return initialize(auth, request, id, params);
  if (sessionId) {
    const info = await loadSession(auth, sessionId);
    if (!info) return sessionMissing();
    const caller = await buildCaller(auth, request, url, "legacy", { id: sessionId, info }, params);
    return respond(caller, request, id, method, params);
  }
  const version = metaOf(params)[META.protocolVersion];
  if (typeof version !== "string") {
    return errorResponse(id, new RpcError(RPC.invalidRequest, "Bad Request: No valid session ID provided"));
  }
  if (version !== MODERN_VERSION) {
    return errorResponse(
      id,
      new RpcError(RPC.unsupportedVersion, "Unsupported protocol version", {
        supported: SUPPORTED_VERSIONS,
        requested: version,
      }),
    );
  }
  const mismatch = checkModernHeaders(request, method, params, version);
  if (mismatch) return errorResponse(id, mismatch);
  const caller = await buildCaller(auth, request, url, "modern", null, params);
  if (method === "subscriptions/listen") {
    try {
      await limitMcpRequest(caller.tokenId ?? caller.viewer.id, null, true);
    } catch (error) {
      if (error instanceof HttpError)
        return errorResponse(id, new RpcError(RPC.invalidRequest, error.message), error.status);
      throw error;
    }
    waitUntil(recordToolUsage(caller.viewer.id, "$subscriptions/listen", false));
    return openListenStream(caller, request, id, params);
  }
  return respond(caller, request, id, method, params);
}

async function handleSessionRequest(auth: Authenticated, request: Request): Promise<Response> {
  // Browser WebSocket clients cannot set headers, so only the upgrade path may carry the id in the query (INT-A30).
  const isUpgrade = request.headers.get("upgrade")?.toLowerCase() === "websocket";
  const sessionId =
    request.headers.get(SESSION_HEADER) ?? (isUpgrade ? new URL(request.url).searchParams.get("session") : null);
  if (!sessionId) return new Response("Method not allowed", { status: 405, headers: { allow: "POST" } });
  if (!(await loadSession(auth, sessionId))) return sessionMissing();
  const hub = sessionHub(sessionId);
  if (request.method === "DELETE") {
    await hub.closeSession();
    return new Response(null, { status: 204 });
  }
  if (request.headers.get("upgrade")?.toLowerCase() === "websocket") {
    const lastEventId =
      request.headers.get("last-event-id") ?? new URL(request.url).searchParams.get("last_event_id") ?? "";
    return hub.fetch(
      new Request(`https://hub/ws?last_event_id=${encodeURIComponent(lastEventId)}`, { headers: request.headers }),
    );
  }
  if (!(request.headers.get("accept") ?? "").includes("text/event-stream")) {
    return new Response("Accept must include text/event-stream", { status: 406 });
  }
  const headers = new Headers();
  const lastEventId = request.headers.get("last-event-id");
  if (lastEventId) headers.set("last-event-id", lastEventId);
  return hub.fetch(new Request("https://hub/sse", { headers, signal: request.signal }));
}

const LOOPBACK = new Set(["localhost", "127.0.0.1", "[::1]"]);

/**
 * Local-development hook that evicts a session's Durable Object, so tests can prove sessions resume from
 * storage after a real eviction. Only exists when both the configured PUBLIC_URL and the request are on
 * loopback (never in a deployment), and only for the caller's own session.
 */
function loopbackDeployment(): boolean {
  try {
    return LOOPBACK.has(new URL(env.PUBLIC_URL).hostname);
  } catch {
    return false;
  }
}

export async function handleTestEviction(request: Request, url: URL): Promise<Response> {
  if (!loopbackDeployment() || !LOOPBACK.has(url.hostname) || request.method !== "POST") {
    return new Response("Not found", { status: 404 });
  }
  const auth = await authenticate(request);
  if (!auth) return unauthorizedResponse(url.origin);
  const sessionId = request.headers.get(SESSION_HEADER) ?? "";
  if (!(await loadSession(auth, sessionId))) return sessionMissing();
  const before = await sessionHub(sessionId).probe();
  await sessionHub(sessionId)
    .evictForTest()
    .catch(() => undefined);
  const after = await sessionHub(sessionId).probe();
  return Response.json({ before, after });
}

/**
 * The `/mcp` endpoint. Speaks 2026-07-28 (stateless, per-request `_meta`) and the 2025 session revisions
 * (`initialize` + `Mcp-Session-Id`, GET stream resumable via `Last-Event-ID`, or a hibernatable WebSocket).
 */
export async function handleMcpEndpoint(request: Request, url: URL): Promise<Response> {
  const origin = request.headers.get("origin");
  if (origin && origin !== url.origin) {
    return errorResponse(null, new RpcError(RPC.invalidRequest, "Origin not allowed"), 403);
  }
  const auth = await authenticate(request);
  if (!auth)
    return unauthorizedResponse(url.origin, request.headers.has("authorization") ? "invalid_token" : undefined);
  if (request.method === "POST") return handlePost(auth, request, url);
  if (request.method === "GET" || request.method === "DELETE") return handleSessionRequest(auth, request);
  return new Response("Method not allowed", { status: 405, headers: { allow: "GET, POST, DELETE" } });
}