import { DurableObject, env } from "cloudflare:workers";
import { loadRepoById } from "../repos";
import { credentialAllows, patCredentialById, tokenIsLive } from "../token-scope";
import { readableEventTypes } from "./event-access";
import { Coordination } from "./coordination";
import type { ClaimRequest, ClaimResult, Lease, LeaseKind, RepoEvent } from "./coordination";
import { SessionLog, sseFrame } from "./hub-session";
import type { SessionInfo } from "./hub-session";
import { META, notification } from "./protocol";
import type { JsonRpcMessage, McpCaller } from "./protocol";

const INFO_KEY = "session";
const SESSION_IDLE_MS = 24 * 60 * 60 * 1000;
const STREAM_MAX_MS = 5 * 60 * 1000;
const KEEPALIVE_MS = 25 * 1000;
const RETRY_MS = 2000;
const SUBSCRIPTION_TTL_MS = SESSION_IDLE_MS;
const MAX_SOCKETS = 4;
const MAX_STREAMS = 4;
const encoder = new TextEncoder();
/** How long a session trusts a granted read check before re-reading the token and repo from D1. */
export const ACCESS_CACHE_MS = 60 * 1000;

type AccessEntry = { at: number; tokenId: string; types: Set<string> };

type Stream = {
  writer: WritableStreamDefaultWriter<Uint8Array>;
  timer: ReturnType<typeof setInterval>;
  deadline: ReturnType<typeof setTimeout> | null;
  onEnd: (() => void) | null;
};

export function sessionHub(sessionId: string) {
  return env.MCP_HUB.get(env.MCP_HUB.idFromName(`session:${sessionId}`));
}

/** Quota key for subscriber slots: always the user, so minting more tokens never buys more slots. */
export function subscriberOwner(caller: Pick<McpCaller, "viewer">): string {
  return caller.viewer.id;
}

export function sessionSubscriber(sessionId: string): string {
  return env.MCP_HUB.idFromName(`session:${sessionId}`).toString();
}

export function repoHub(repoId: string) {
  return env.MCP_HUB.get(env.MCP_HUB.idFromName(`repo:${repoId}`));
}

export function eventJson(event: RepoEvent) {
  return { seq: event.seq, type: event.type, at: event.at, data: JSON.parse(event.data) as unknown };
}

export function eventsUri(fullName: string): string {
  return `gitfor-events://${fullName}`;
}

/**
 * One class, two roles keyed by name: `session:<id>` holds an MCP session (replay log, standalone
 * SSE stream, hibernatable WebSocket) and `repo:<id>` coordinates agents on one repository (leases,
 * subscriptions, recent events). Idle sessions hold no connection and cost nothing until the next call.
 */
export class McpHub extends DurableObject<Env> {
  private coordinationState: Coordination | null = null;
  private logState: SessionLog | null = null;
  private readonly streams = new Set<Stream>();
  /** Granted (never denied) repo access per repo; token liveness is re-checked on every hit. */
  private readonly access = new Map<string, AccessEntry>();
  /** Changes whenever the runtime constructs a fresh instance, e.g. after eviction. */
  private readonly instanceId = crypto.randomUUID();

  private get coordination(): Coordination {
    this.coordinationState ??= new Coordination(this.ctx.storage.sql);
    return this.coordinationState;
  }

  private get log(): SessionLog {
    this.logState ??= new SessionLog(this.ctx.storage.sql);
    return this.logState;
  }

  async openSession(info: SessionInfo): Promise<void> {
    await this.ctx.storage.put(INFO_KEY, info);
    await this.touch();
  }

  async session(): Promise<SessionInfo | null> {
    const info = (await this.ctx.storage.get<SessionInfo>(INFO_KEY)) ?? null;
    if (info) await this.touch();
    return info;
  }

  async updateSession(
    patch: Partial<Pick<SessionInfo, "resourceUris" | "repos" | "toolEvents">>,
  ): Promise<SessionInfo | null> {
    const info = await this.ctx.storage.get<SessionInfo>(INFO_KEY);
    if (!info) return null;
    const next = { ...info, ...patch };
    await this.ctx.storage.put(INFO_KEY, next);
    return next;
  }

  async closeSession(): Promise<void> {
    const info = await this.ctx.storage.get<SessionInfo>(INFO_KEY);
    for (const stream of this.streams) this.endStream(stream);
    for (const socket of this.ctx.getWebSockets()) socket.close(1000, "Session closed");
    if (info) await Promise.allSettled(info.repos.map((repoId) => repoHub(repoId).unsubscribe(this.subscriberName())));
    await this.ctx.storage.deleteAll();
    this.logState = null;
  }

  /** Test hook: identifies the live instance and its in-memory streams. */
  probe(): { instance: string; streams: number } {
    return { instance: this.instanceId, streams: this.streams.size };
  }

  /**
   * Test hook: discards this instance exactly as an eviction would, dropping all in-memory state
   * (open streams, timers, caches). Storage, alarms and the replay log survive.
   */
  evictForTest(): void {
    this.ctx.abort("test eviction");
  }

  /** Repo hubs address subscribers by object id, so a session never has to know its own name. */
  subscriberName(): string {
    return this.ctx.id.toString();
  }

  private async touch(): Promise<void> {
    await this.ctx.storage.setAlarm(Date.now() + SESSION_IDLE_MS);
  }

  override async alarm(): Promise<void> {
    if (await this.ctx.storage.get(INFO_KEY)) await this.closeSession();
  }

  /** Called by repo hubs. Re-checks the token's read access so revoked grants stop receiving events. */
  async deliver(repoId: string, fullName: string, event: RepoEvent): Promise<boolean> {
    const info = await this.ctx.storage.get<SessionInfo>(INFO_KEY);
    if (!info) return false;
    const types = await this.readableTypes(info.tokenId, repoId);
    if (!types) {
      // Events logged while access was live must not replay to a reconnect after it was revoked.
      this.log.clearRepo(repoId);
      return false;
    }
    if (!types.has(event.type)) return true;
    const uri = eventsUri(fullName);
    if (info.kind === "listen") {
      if (!info.resourceUris.includes(uri)) return true;
      this.broadcast(
        notification("notifications/resources/updated", { uri, _meta: { [META.subscriptionId]: info.subscriptionId } }),
        false,
      );
      return true;
    }
    if (info.resourceUris.includes(uri))
      this.broadcast(notification("notifications/resources/updated", { uri }), true, repoId);
    const wanted = info.toolEvents[repoId];
    if (wanted && (wanted.includes("*") || wanted.includes(event.type))) {
      this.broadcast(
        notification("notifications/gitfor/event", { repo: fullName, uri, event: eventJson(event) }),
        true,
        repoId,
      );
    }
    return true;
  }

  private async readableTypes(tokenId: string, repoId: string): Promise<Set<string> | null> {
    const cached = this.access.get(repoId);
    if (cached && cached.tokenId === tokenId && Date.now() - cached.at < ACCESS_CACHE_MS) {
      // Repo access may be cached, token liveness never is: a revoked or expired token fails closed at once.
      if (await tokenIsLive(tokenId)) return cached.types;
      this.access.clear();
      return null;
    }
    this.access.delete(repoId);
    const [credential, repo] = await Promise.all([patCredentialById(tokenId), loadRepoById(repoId)]);
    if (!credential?.viewer || !repo || !(await credentialAllows(repo, credential, "contents", "read"))) return null;
    const types = await readableEventTypes(repo, credential);
    this.access.set(repoId, { at: Date.now(), tokenId, types });
    return types;
  }

  /**
   * Forgets cached access checks so the next delivery re-reads the token and repo (after a revocation). That
   * repository's replay backlog goes too, so a reconnect can't replay events it may no longer read; other
   * repositories' backlog stays.
   */
  async invalidateAccess(repoId?: string): Promise<void> {
    if (repoId) {
      this.access.delete(repoId);
      this.log.clearRepo(repoId);
    } else {
      this.access.clear();
      this.log.clear();
    }
  }

  /** Repo role: tells every subscribed session to drop its cached access to this repo. */
  async invalidateSubscribers(repoId: string): Promise<void> {
    const names = this.coordination.allSubscribers();
    await Promise.allSettled(
      names.map((name) => this.env.MCP_HUB.get(this.env.MCP_HUB.idFromString(name)).invalidateAccess(repoId)),
    );
  }

  private broadcast(message: JsonRpcMessage, persist: boolean, repoId: string | null = null): void {
    const entry = persist ? this.log.append(message, repoId) : null;
    const frame = encoder.encode(sseFrame(entry?.data ?? JSON.stringify(message), entry?.id));
    for (const stream of this.streams) stream.writer.write(frame).catch(() => this.endStream(stream));
    const payload = JSON.stringify(entry ? { id: entry.id, message } : { message });
    for (const socket of this.ctx.getWebSockets()) {
      try {
        socket.send(payload);
      } catch {
        socket.close(1011, "Send failed");
      }
    }
  }

  private endStream(stream: Stream): void {
    if (!this.streams.delete(stream)) return;
    clearInterval(stream.timer);
    if (stream.deadline) clearTimeout(stream.deadline);
    stream.writer.close().catch(() => undefined);
    stream.onEnd?.();
  }

  private openStream(
    preamble: string[],
    options: { finalFrame?: string; onEnd?: () => void; signal?: AbortSignal } = {},
  ): Response {
    const oldest = this.streams.size >= MAX_STREAMS ? this.streams.values().next().value : undefined;
    if (oldest) this.endStream(oldest);
    const { readable, writable } = new TransformStream<Uint8Array, Uint8Array>();
    const writer = writable.getWriter();
    const stream: Stream = {
      writer,
      timer: setInterval(() => writer.write(encoder.encode(":\n\n")).catch(() => this.endStream(stream)), KEEPALIVE_MS),
      deadline: null,
      onEnd: options.onEnd ?? null,
    };
    this.streams.add(stream);
    options.signal?.addEventListener("abort", () => this.endStream(stream));
    for (const chunk of preamble) void writer.write(encoder.encode(chunk));
    stream.deadline = setTimeout(() => {
      if (options.finalFrame && this.streams.has(stream)) void writer.write(encoder.encode(options.finalFrame));
      this.endStream(stream);
    }, STREAM_MAX_MS);
    return new Response(readable, {
      headers: {
        "content-type": "text/event-stream",
        "cache-control": "no-cache, no-transform",
        "x-accel-buffering": "no",
      },
    });
  }

  /** Replay only while the session's token is still live; a revoked token's backlog is dropped. */
  private async replay(lastEventId: string | null): Promise<Array<{ id: string; data: string }>> {
    const info = await this.ctx.storage.get<SessionInfo>(INFO_KEY);
    if (!info || !(await patCredentialById(info.tokenId))?.viewer) {
      this.log.clear();
      return [];
    }
    return this.log.since(lastEventId);
  }

  /** `/sse`: legacy standalone stream with replay; `/ws`: hibernatable socket; `/listen`: 2026 subscriptions. */
  override async fetch(request: Request): Promise<Response> {
    const info = await this.ctx.storage.get<SessionInfo>(INFO_KEY);
    if (!info) return new Response("Session not found", { status: 404 });
    await this.touch();
    const url = new URL(request.url);
    if (url.pathname === "/ws") {
      if (request.headers.get("upgrade") !== "websocket") return new Response("Expected websocket", { status: 426 });
      if (this.ctx.getWebSockets().length >= MAX_SOCKETS) return new Response("Too many connections", { status: 429 });
      const pair = new WebSocketPair();
      this.ctx.acceptWebSocket(pair[1]);
      for (const entry of await this.replay(url.searchParams.get("last_event_id") ?? "get-0")) {
        pair[1].send(JSON.stringify({ id: entry.id, message: JSON.parse(entry.data) }));
      }
      return new Response(null, { status: 101, webSocket: pair[0] });
    }
    if (url.pathname === "/sse") {
      const replay = await this.replay(request.headers.get("last-event-id"));
      return this.openStream([`retry: ${RETRY_MS}\n\n`, ...replay.map((entry) => sseFrame(entry.data, entry.id))], {
        signal: request.signal,
      });
    }
    if (url.pathname === "/listen") {
      const ack = notification("notifications/subscriptions/acknowledged", {
        _meta: { [META.subscriptionId]: info.subscriptionId },
        notifications: { resourceSubscriptions: info.resourceUris },
      });
      const done = {
        jsonrpc: "2.0",
        id: info.subscriptionId,
        result: { resultType: "complete", _meta: { [META.subscriptionId]: info.subscriptionId } },
      };
      const response = this.openStream([sseFrame(JSON.stringify(ack))], {
        finalFrame: sseFrame(JSON.stringify(done)),
        onEnd: () => void this.closeSession(),
        signal: request.signal,
      });
      return response;
    }
    return new Response("Not found", { status: 404 });
  }

  async webSocketMessage(socket: WebSocket, message: string | ArrayBuffer): Promise<void> {
    if (message === "ping") socket.send("pong");
  }

  async webSocketClose(socket: WebSocket): Promise<void> {
    socket.close();
  }

  claim(request: ClaimRequest): ClaimResult {
    return this.coordination.claim(request);
  }

  release(holder: string, ids: string[] | null): number {
    return this.coordination.release(holder, ids);
  }

  leases(kind?: LeaseKind): Lease[] {
    return this.coordination.leases(kind);
  }

  subscribe(name: string, events: string[], owner: string | null = null): number {
    return this.coordination.subscribe(name, events, SUBSCRIPTION_TTL_MS, owner);
  }

  unsubscribe(name: string): void {
    this.coordination.unsubscribe(name);
  }

  eventsSince(after: number, limit: number, types: string[] | null): { events: RepoEvent[]; latest: number } {
    return { events: this.coordination.eventsSince(after, limit, types), latest: this.coordination.latestSeq() };
  }

  /** Records a repository event and fans it out to every subscribed session hub. */
  async publish(repoId: string, fullName: string, type: string, data: string): Promise<number> {
    const event = this.coordination.record(type, data);
    const names = this.coordination.subscribersFor(type);
    const results = await Promise.allSettled(
      names.map((name) => this.env.MCP_HUB.get(this.env.MCP_HUB.idFromString(name)).deliver(repoId, fullName, event)),
    );
    results.forEach((result, index) => {
      if (result.status === "fulfilled" && !result.value) this.coordination.unsubscribe(names[index]);
    });
    return event.seq;
  }
}