// ---------------------------------------------------------------------------
// MCP Session Durable Object — holds MCP server + engine per session
// ---------------------------------------------------------------------------

import { DurableObject, env } from "cloudflare:workers";
import { Data, Effect, Layer } from "effect";
import * as OtelTracer from "@effect/opentelemetry/Tracer";
import * as Sentry from "@sentry/cloudflare";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { WorkerTransport, type TransportState } from "agents/mcp";
import { drizzle } from "drizzle-orm/postgres-js";
import postgres from "postgres";

import { createExecutorMcpServer } from "@executor/host-mcp";
import { buildExecuteDescription } from "@executor/execution";
import type { DrizzleDb, DbServiceShape } from "./services/db";

// Import directly from core-shared-services, NOT from ./api/layers.ts.
// The full layers module pulls in `auth/handlers.ts` → `@tanstack/react-start/server`,
// which uses a `#tanstack-start-entry` subpath specifier that breaks module
// load under vitest-pool-workers. The DO only needs the core two services
// (WorkOSAuth + AutumnService), so we import them from the tight module.
import { CoreSharedServices } from "./api/core-shared-services";
import { UserStoreService } from "./auth/context";
import { resolveOrganization } from "./auth/resolve-organization";
import { server } from "./env";
import { DbService, combinedSchema } from "./services/db";
import { makeExecutionStack } from "./services/execution-stack";
import { DoTelemetryLive } from "./services/telemetry";

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

export type McpSessionInit = {
  organizationId: string;
};

const HEARTBEAT_MS = 30 * 1000;
const SESSION_TIMEOUT_MS = 5 * 60 * 1000;
const TRANSPORT_STATE_KEY = "transport";
const SESSION_META_KEY = "session-meta";

// ---------------------------------------------------------------------------
// Errors
// ---------------------------------------------------------------------------

class OrganizationNotFoundError extends Data.TaggedError("OrganizationNotFoundError")<{
  readonly organizationId: string;
}> {}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

const jsonRpcError = (status: number, code: number, message: string) =>
  new Response(JSON.stringify({ jsonrpc: "2.0", error: { code, message }, id: null }), {
    status,
    headers: { "content-type": "application/json" },
  });

// W3C traceparent propagation across the worker→DO boundary. mcp.ts injects
// the worker-side `mcp.request` SpanContext as a `traceparent` header on
// forwarded requests (and as a second arg to `init()`). We parse it here and
// use `OtelTracer.withSpanContext` to stitch the DO's root span under the
// worker span so the entire logical request lives in one Axiom trace.
const TRACEPARENT_PATTERN = /^([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$/;

type IncomingSpanContext = {
  readonly traceId: string;
  readonly spanId: string;
  readonly traceFlags: number;
};

const parseTraceparent = (value: string | null | undefined): IncomingSpanContext | null => {
  if (!value) return null;
  const match = TRACEPARENT_PATTERN.exec(value);
  if (!match) return null;
  return { traceId: match[2]!, spanId: match[3]!, traceFlags: parseInt(match[4]!, 16) };
};

const withIncomingParent = <A, E, R>(
  traceparent: string | null | undefined,
  effect: Effect.Effect<A, E, R>,
): Effect.Effect<A, E, R> => {
  const parsed = parseTraceparent(traceparent);
  return parsed ? OtelTracer.withSpanContext(effect, parsed) : effect;
};

type DbHandle = DbServiceShape & { end: () => Promise<void> };
type SessionMeta = {
  readonly organizationId: string;
  readonly organizationName: string;
};

/**
 * Base DB handle factory for MCP session runtimes.
 *
 * The production DO keeps one postgres.js socket for the session lifetime.
 * The workerd test pool rejects that socket on the next request because the
 * underlying I/O object is request-bound. Tests therefore opt into a
 * request-scoped runtime and rebuild the server + DB handle per POST/DELETE
 * while keeping only the MCP transport state in DO storage.
 */
const makeDbHandle = (options: {
  readonly idleTimeout: number;
  readonly maxLifetime: number;
}): DbHandle => {
  const connectionString = env.HYPERDRIVE?.connectionString ?? server.DATABASE_URL;
  const sql = postgres(connectionString, {
    max: 1,
    idle_timeout: options.idleTimeout,
    max_lifetime: options.maxLifetime,
    connect_timeout: 10,
    onnotice: () => undefined,
  });
  return {
    sql,
    db: drizzle(sql, { schema: combinedSchema }) as DrizzleDb,
    end: () => sql.end({ timeout: 0 }).catch(() => undefined),
  };
};

const makeLongLivedDb = (): DbHandle => makeDbHandle({ idleTimeout: 20, maxLifetime: 300 });

const makeRequestScopedDb = (): DbHandle => makeDbHandle({ idleTimeout: 0, maxLifetime: 60 });

const makeResolveOrganizationServices = (dbHandle: DbHandle) => {
  const DbLive = Layer.succeed(DbService, { sql: dbHandle.sql, db: dbHandle.db });
  const UserStoreLive = UserStoreService.Live.pipe(Layer.provide(DbLive));
  return Layer.mergeAll(DbLive, UserStoreLive, CoreSharedServices);
};

// Session services DON'T re-provide `DoTelemetryLive` — that would install a
// second WebSdk tracer in the nested Effect scope, disconnecting every
// child span from the outer `McpSessionDO.init` / `McpSessionDO.handleRequest`
// trace. Tracer comes from the outermost `Effect.provide(DoTelemetryLive)`
// at the DO method boundary.
const makeSessionServices = (dbHandle: DbHandle) =>
  makeResolveOrganizationServices(dbHandle);

const resolveSessionMeta = Effect.fn("McpSessionDO.resolveSessionMeta")(function* (
  organizationId: string,
) {
    const org = yield* resolveOrganization(organizationId);
    if (!org) {
      return yield* new OrganizationNotFoundError({ organizationId });
    }
    return {
      organizationId: org.id,
      organizationName: org.name,
    } satisfies SessionMeta;
  });

const requestScopedRuntimeEnabled = server.MCP_SESSION_REQUEST_SCOPED_RUNTIME === "true";

// ---------------------------------------------------------------------------
// Durable Object
// ---------------------------------------------------------------------------

export class McpSessionDO extends DurableObject {
  private mcpServer: McpServer | null = null;
  private transport: WorkerTransport | null = null;
  private initialized = false;
  private lastActivityMs = 0;
  private dbHandle: DbHandle | null = null;
  private sessionMeta: SessionMeta | null = null;

  private makeStorage() {
    return {
      get: async (): Promise<TransportState | undefined> => {
        return await this.ctx.storage.get<TransportState>(TRANSPORT_STATE_KEY);
      },
      set: async (state: TransportState): Promise<void> => {
        await this.ctx.storage.put(TRANSPORT_STATE_KEY, state);
      },
    };
  }

  private async loadSessionMeta(): Promise<SessionMeta | null> {
    if (this.sessionMeta) return this.sessionMeta;
    const stored = await this.ctx.storage.get<SessionMeta>(SESSION_META_KEY);
    this.sessionMeta = stored ?? null;
    return this.sessionMeta;
  }

  private async saveSessionMeta(sessionMeta: SessionMeta): Promise<void> {
    this.sessionMeta = sessionMeta;
    await this.ctx.storage.put(SESSION_META_KEY, sessionMeta);
  }

  private async clearSessionState(): Promise<void> {
    this.sessionMeta = null;
    this.initialized = false;
    this.lastActivityMs = 0;

    await Promise.all([
      this.ctx.storage.delete(TRANSPORT_STATE_KEY).catch(() => false),
      this.ctx.storage.delete(SESSION_META_KEY).catch(() => false),
    ]);
  }

  private createConnectedRuntimeEffect(
    sessionMeta: SessionMeta,
    options: { readonly dbHandle: DbHandle; readonly enableJsonResponse?: boolean },
  ) {
    const self = this;
    return Effect.gen(function* () {
      const { executor, engine } = yield* makeExecutionStack(
        sessionMeta.organizationId,
        sessionMeta.organizationName,
      );
      // Build the description here so the two postgres queries it runs
      // (`executor.sources.list` + `executor.tools.list`) land as
      // children of `McpSessionDO.createRuntime`. host-mcp would
      // otherwise call `Effect.runPromise(engine.getDescription)` at
      // its async MCP-SDK boundary and orphan those sub-spans.
      const description = yield* buildExecuteDescription(executor);
      const mcpServer = yield* Effect.promise(() =>
        createExecutorMcpServer({ engine, description }),
      ).pipe(Effect.withSpan("McpSessionDO.createExecutorMcpServer"));
      const transport = new WorkerTransport({
        sessionIdGenerator: () => self.ctx.id.toString(),
        storage: self.makeStorage(),
        enableJsonResponse: options.enableJsonResponse,
      });
      yield* Effect.promise(() => mcpServer.connect(transport)).pipe(
        Effect.withSpan("McpSessionDO.transport.connect"),
      );
      return { mcpServer, transport };
    }).pipe(
      Effect.withSpan("McpSessionDO.createRuntime"),
      Effect.provide(makeSessionServices(options.dbHandle)),
    );
  }

  private resolveAndStoreSessionMetaEffect(token: McpSessionInit) {
    const self = this;
    return Effect.gen(function* () {
      const dbHandle = makeRequestScopedDb();
      try {
        const sessionMeta = yield* resolveSessionMeta(token.organizationId).pipe(
          Effect.provide(makeResolveOrganizationServices(dbHandle)),
        );
        yield* Effect.promise(() => self.saveSessionMeta(sessionMeta));
        return sessionMeta;
      } finally {
        yield* Effect.promise(() => dbHandle.end());
      }
    });
  }

  async init(token: McpSessionInit, traceparent?: string): Promise<void> {
    if (this.initialized) return;
    return Effect.runPromise(
      this.doInitEffect(token).pipe(
        Effect.withSpan("McpSessionDO.init", {
          attributes: { "mcp.auth.organization_id": token.organizationId },
        }),
        (eff) => withIncomingParent(traceparent, eff),
        Effect.provide(DoTelemetryLive),
      ),
    );
  }

  private doInitEffect(token: McpSessionInit) {
    const self = this;
    // Single Effect chain so every sub-span (resolveSessionMeta,
    // createRuntime, createScopedExecutor, createExecutorMcpServer,
    // transport.connect, storage.setAlarm) lands as a child of
    // `McpSessionDO.init`. The prior implementation called
    // `Effect.runPromise` nested inside an async function, which orphaned
    // each sub-span into its own root trace and made init opaque —
    // dashboard saw one 2.77s span with nothing under it.
    return Effect.gen(function* () {
      const sessionMeta = yield* self.resolveAndStoreSessionMetaEffect(token);

      if (!requestScopedRuntimeEnabled) {
        self.dbHandle = makeLongLivedDb();
        const runtime = yield* self.createConnectedRuntimeEffect(sessionMeta, {
          dbHandle: self.dbHandle,
        });
        self.mcpServer = runtime.mcpServer;
        self.transport = runtime.transport;
      }

      self.initialized = true;
      self.lastActivityMs = Date.now();

      yield* Effect.promise(() => self.ctx.storage.setAlarm(Date.now() + HEARTBEAT_MS)).pipe(
        Effect.withSpan("McpSessionDO.setAlarm"),
      );
    }).pipe(
      Effect.tapErrorCause((cause) =>
        Effect.sync(() => {
          console.error("[mcp-session] init failed:", cause);
        }),
      ),
      Effect.catchAllCause((cause) =>
        Effect.gen(function* () {
          yield* Effect.promise(() => self.cleanup());
          return yield* Effect.failCause(cause);
        }),
      ),
      Effect.orDie,
    );
  }

  private async handleRequestWithRequestScopedRuntime(request: Request): Promise<Response> {
    const sessionMeta = await this.loadSessionMeta();
    if (!sessionMeta) {
      return jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect");
    }

    this.initialized = true;
    this.lastActivityMs = Date.now();

    let dbHandle: DbHandle | null = null;
    let mcpServer: McpServer | null = null;
    let transport: WorkerTransport | null = null;

    try {
      dbHandle = makeRequestScopedDb();
      const runtime = await Effect.runPromise(
        this.createConnectedRuntimeEffect(sessionMeta, {
          dbHandle,
          enableJsonResponse: request.method !== "GET",
        }).pipe(Effect.provide(DoTelemetryLive)),
      );
      mcpServer = runtime.mcpServer;
      transport = runtime.transport;

      const response = await transport.handleRequest(request);
      if (request.method === "DELETE") {
        await this.clearSessionState();
      }
      return response;
    } catch (err) {
      console.error(
        "[mcp-session] request-scoped handleRequest error:",
        err instanceof Error ? err.stack : err,
      );
      Sentry.captureException(err);
      return jsonRpcError(500, -32603, "Internal error");
    } finally {
      await transport?.close().catch(() => undefined);
      await mcpServer?.close().catch(() => undefined);
      await dbHandle?.end();
    }
  }

  async handleRequest(request: Request): Promise<Response> {
    // Wrap the dispatch in an Effect span so every DO request — not just
    // the rare new-session `init()` — shows up in Axiom. Basic attributes
    // only (method, session-id presence, response status); rich client
    // fingerprint stays on the edge `mcp.request` span, which shares a
    // trace_id with this one.
    const traceparent = request.headers.get("traceparent");
    const program = Effect.promise(() => this.dispatchRequest(request)).pipe(
      Effect.tap((response) =>
        Effect.annotateCurrentSpan({
          "mcp.response.status_code": response.status,
        }),
      ),
      Effect.withSpan("McpSessionDO.handleRequest", {
        attributes: {
          "mcp.request.method": request.method,
          "mcp.request.session_id_present": !!request.headers.get("mcp-session-id"),
        },
      }),
      (eff) => withIncomingParent(traceparent, eff),
      Effect.provide(DoTelemetryLive),
    );
    return Effect.runPromise(program);
  }

  private async dispatchRequest(request: Request): Promise<Response> {
    if (requestScopedRuntimeEnabled) {
      return this.handleRequestWithRequestScopedRuntime(request);
    }

    if (!this.initialized || !this.transport) {
      return jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect");
    }

    this.lastActivityMs = Date.now();

    try {
      const response = await this.transport.handleRequest(request);
      if (request.method === "DELETE") {
        await this.cleanup();
      }
      return response;
    } catch (err) {
      console.error("[mcp-session] handleRequest error:", err instanceof Error ? err.stack : err);
      Sentry.captureException(err);
      return jsonRpcError(500, -32603, "Internal error");
    }
  }

  async alarm(): Promise<void> {
    const program = Effect.promise(() => this.runAlarm()).pipe(
      Effect.withSpan("McpSessionDO.alarm"),
      Effect.provide(DoTelemetryLive),
    );
    return Effect.runPromise(program);
  }

  private async runAlarm(): Promise<void> {
    const idleMs = Date.now() - this.lastActivityMs;
    if (idleMs >= SESSION_TIMEOUT_MS) {
      await this.cleanup();
      return;
    }
    await this.ctx.storage.setAlarm(Date.now() + HEARTBEAT_MS);
  }

  private async cleanup(): Promise<void> {
    if (this.transport) {
      await this.transport.close().catch(() => undefined);
      this.transport = null;
    }
    if (this.mcpServer) {
      await this.mcpServer.close().catch(() => undefined);
      this.mcpServer = null;
    }
    if (this.dbHandle) {
      await this.dbHandle.end();
      this.dbHandle = null;
    }
    await this.clearSessionState();
  }
}