import { getConfig, type AppConfig } from "../config.js";
import type { AgentSession } from "../agents/manager.js";
import type { ProvisionedResources } from "../agents/provisioning.js";
import {
  getAgentManager,
  getResourceProvisioningOrchestrator,
} from "../agents/runtime.js";
import {
  getDefaultMcpRegistry,
  loadMcpConfigFromEnv,
  type LoadedMcpConfig,
  type McpRegistry,
} from "../mcp.js";

export type SessionProvider = "codex" | "claude";

export type SessionContext = {
  provider: SessionProvider;
  sessionId: string;
  mcp: LoadedMcpConfig;
  promptWithContext: string;
  runtimeInjection: Awaited<ReturnType<ProvisionedResources["buildRuntimeInjection"]>>;
  runInSession: <T>(run: () => Promise<T>) => Promise<T>;
  close: () => Promise<void>;
};

export async function createSessionContext(
  provider: SessionProvider,
  input: {
    prompt: string;
    config?: Readonly<AppConfig>;
    mcpRegistry?: McpRegistry;
  },
): Promise<SessionContext> {
  const config = input.config ?? getConfig();
  const mcpRegistry = input.mcpRegistry ?? getDefaultMcpRegistry();
  const agentManager = getAgentManager(config);
  const agentSession = agentManager.createSession();
  const resourceProvisioning = getResourceProvisioningOrchestrator(config);

  let provisionedResources: ProvisionedResources | undefined;
  let closed = false;

  const close = async (): Promise<void> => {
    if (closed) {
      return;
    }
    closed = true;

    if (provisionedResources) {
      await provisionedResources.release();
    }

    agentSession.close();
  };

  try {
    provisionedResources = await resourceProvisioning.provisionSession({
      sessionId: agentSession.id,
      resources: [{ kind: "git-worktree" }, { kind: "port-range" }],
    });

    const runtimeInjection = await provisionedResources.buildRuntimeInjection({
      discoveryFileRelativePath: config.discovery.fileRelativePath,
      baseEnv: process.env,
    });

    const promptWithContext = provisionedResources.composePrompt(input.prompt, [
      `Discovery file: ${runtimeInjection.discoveryFilePath}`,
      "Resource env vars are pre-injected (AGENT_WORKTREE_PATH, AGENT_PORT_RANGE_START, AGENT_PORT_RANGE_END, AGENT_PORT_PRIMARY).",
    ]);

    const mcp = loadMcpConfigFromEnv(
      {
        providerHint: provider,
        prompt: input.prompt,
      },
      {
        config,
        registry: mcpRegistry,
      },
    );

    return {
      provider,
      sessionId: agentSession.id,
      mcp,
      promptWithContext,
      runtimeInjection,
      runInSession: <T>(run: () => Promise<T>) =>
        runWithAgentSession(agentSession, run),
      close,
    };
  } catch (error) {
    await close();
    throw error;
  }
}

async function runWithAgentSession<T>(agentSession: AgentSession, run: () => Promise<T>): Promise<T> {
  return agentSession.runAgent({
    depth: 0,
    run,
  });
}