import { randomUUID } from "node:crypto";
import { mkdir, realpath, rm } from "node:fs/promises";
import path from "node:path";
import type { ImageContent } from "@earendil-works/pi-ai";
import {
  SessionManager,
  SettingsManager,
  createAgentSession,
  isToolCallEventType,
  type AgentSession,
  type AgentSessionEvent,
  type InlineExtension,
} from "@earendil-works/pi-coding-agent";
import type { ModelConnectionInput, SendMessageInput } from "@agent-studio/shared";
import { InMemoryApprovalBroker, type ApprovalBroker } from "./approval.js";
import type { FileAttachmentStore } from "./attachments.js";
import { HarnessError, redactValue, redactServerSecret, safeErrorMessage } from "./errors.js";
import type { EventStore } from "./event-store.js";
import type { ModelConnectionStore } from "./model-store.js";
import { mapPiEvent } from "./pi-event-adapter.js";
import { SafeWorkspace } from "./safe-workspace.js";
import type { SkillCatalog } from "./skills.js";
import type { SkillIndexEntry } from "./skills.js";
import { SwadsGateway } from "./swads-gateway.js";
import { SwadsReportTools } from "./swads-tools.js";
import { ReportStore } from "./report-store.js";
import { createMockAdsMetricsTool, createSaveReportDraftTool } from "./tools.js";

export interface AgentSessionPort {
  readonly modelId: string;
  readonly supportsImages: boolean;
  prompt(text: string, options?: { images?: unknown[] }): Promise<void>;
  abort(): Promise<void>;
  subscribe(listener: (event: unknown) => void): () => void;
  dispose(): void;
}

export interface AgentSessionFactoryInput {
  sessionId: string;
  cwd: string;
  workspace: SafeWorkspace;
  getRun: () => { runId: string; signal: AbortSignal; assistantMessageId: string } | undefined;
  skillRequested: (skill: SkillIndexEntry, toolCallId: string) => void;
  skillLoaded: (skill: SkillIndexEntry, toolCallId: string) => void;
  toolRequested: (toolCallId: string, toolName: string, argumentsValue: unknown) => void;
}
export interface AgentSessionFactory { create(input: AgentSessionFactoryInput): Promise<AgentSessionPort> }

class PiAgentSessionPort implements AgentSessionPort {
  readonly modelId: string;
  readonly supportsImages: boolean;
  constructor(private readonly session: AgentSession) {
    if (!session.model) throw new Error("Pi AgentSession was created without an explicit model");
    this.modelId = session.model.id;
    this.supportsImages = session.model.input.includes("image");
  }
  async prompt(text: string, options?: { images?: unknown[] }): Promise<void> {
    const images = options?.images as ImageContent[] | undefined;
    await this.session.prompt(text, images ? { images } : undefined);
  }
  abort(): Promise<void> { return this.session.abort(); }
  subscribe(listener: (event: unknown) => void): () => void { return this.session.subscribe(listener); }
  dispose(): void { this.session.dispose(); }
}

export class PiAgentSessionFactory implements AgentSessionFactory {
  constructor(
    private readonly fixturePath: string,
    private readonly catalog: SkillCatalog,
    private readonly models: ModelConnectionStore,
    private readonly approvals: ApprovalBroker,
    private readonly gateway = new SwadsGateway(),
    private readonly reports = new ReportStore(path.resolve(fixturePath, "../../workspace/reports"), path.resolve(fixturePath, "../../.agents/skills/swads-daily-report/assets/report.css")),
  ) {}

  async create(input: AgentSessionFactoryInput): Promise<AgentSessionPort> {
    const resolved = await this.models.createRuntime(input.sessionId);
    const swads = new SwadsReportTools(this.gateway, this.reports, input.getRun);
    const extension: InlineExtension = { name: "agent-studio-safety", hidden: true, factory: (pi) => {
      const pendingSkills = new Map<string, SkillIndexEntry>();
      pi.on("tool_call", async (event) => {
        const run = input.getRun();
        if (!run) return { block: true, reason: "No active Harness run" };
        const known = new Set(["read", "grep", "find", "ls", "mock_ads_metrics", "save_report_draft", "swads_cli", "render_swads_daily_report"]);
        if (!known.has(event.toolName)) {
          input.toolRequested(event.toolCallId, event.toolName, publicToolArguments(event.input, input.cwd));
          return { block: true, reason: "Tool is not allowed by the Harness policy", terminate: true };
        }
        let matchedSkill: SkillIndexEntry | undefined;
        if (isToolCallEventType("read", event)) {
          const skill = await this.catalog.matchReadPath(event.input.path, input.cwd);
          if (skill) {
            matchedSkill = skill;
            pendingSkills.set(event.toolCallId, skill);
            input.skillRequested(skill, event.toolCallId);
          } else await input.workspace.assertReadPath(event.input.path);
        } else if (["grep", "find", "ls"].includes(event.toolName)) {
          const maybePath = (event.input as Record<string, unknown>).path;
          if (typeof maybePath === "string") await input.workspace.assertReadPath(maybePath);
        }
        input.toolRequested(event.toolCallId, event.toolName, event.toolName === "render_swads_daily_report" ? { status: "requested" } : publicToolArguments(event.input, input.cwd, matchedSkill));
        if (event.toolName === "save_report_draft") {
          const decision = await this.approvals.request({ sessionId: input.sessionId, runId: run.runId, toolCallId: event.toolCallId, toolName: event.toolName, arguments: event.input, riskLevel: "medium" }, run.signal);
          if (decision.status !== "approved") return { block: true, reason: `Human approval ${decision.status}` };
        }
        return undefined;
      });
      pi.on("tool_result", async (event) => {
        const run = input.getRun();
        const skill = pendingSkills.get(event.toolCallId);
        if (run && skill && !event.isError && event.toolName === "read") {
          input.skillLoaded(skill, event.toolCallId);
          pendingSkills.delete(event.toolCallId);
        }
      });
    } };
    const loader = await this.catalog.createSessionLoader(input.cwd, path.join(input.cwd, ".pi-agent"), [extension]);
    const { session } = await createAgentSession({
      cwd: input.cwd,
      agentDir: path.join(input.cwd, ".pi-agent"),
      modelRuntime: resolved.runtime,
      model: resolved.model,
      tools: ["read", "grep", "find", "ls", "mock_ads_metrics", "save_report_draft", "swads_cli", "render_swads_daily_report"],
      customTools: [createMockAdsMetricsTool(this.fixturePath), createSaveReportDraftTool(input.workspace), ...swads.definitions()],
      resourceLoader: loader,
      sessionManager: SessionManager.inMemory(input.cwd),
      settingsManager: SettingsManager.inMemory(),
    });
    return new PiAgentSessionPort(session);
  }
}

export interface HarnessSession {
  sessionId: string;
  cwd: string;
  workspace: SafeWorkspace;
  busy: boolean;
  currentRun: { runId: string; controller: AbortController; assistantMessageId: string; cancelReason?: "user" | "timeout" | "session_deleted" | "shutdown" } | undefined;
}
export interface CreateSessionInput { id?: string }
export interface SessionRegistry {
  create(input?: CreateSessionInput): Promise<HarnessSession>;
  get(sessionId: string): HarnessSession | undefined;
  delete(sessionId: string): Promise<void>;
}

interface InternalSession extends HarnessSession { agent: AgentSessionPort | undefined; unsubscribe: (() => void) | undefined; timeout: ReturnType<typeof setTimeout> | undefined; disposed: boolean; requestedToolCalls: Set<string> }

export class InMemorySessionRegistry implements SessionRegistry {
  private readonly sessions = new Map<string, InternalSession>();
  readonly approvals: ApprovalBroker;
  constructor(
    private readonly sessionsRoot: string,
    private readonly events: EventStore,
    private readonly attachments: FileAttachmentStore,
    private readonly models: ModelConnectionStore,
    private readonly catalog: SkillCatalog,
    private factory: AgentSessionFactory | undefined,
    approvalTimeoutMs = 60_000,
    private readonly runTimeoutMs = 120_000,
  ) {
    this.approvals = new InMemoryApprovalBroker(approvalTimeoutMs, (resolution) => {
      this.events.append(resolution.request.sessionId, { type: "approval.resolved", runId: resolution.request.runId, payload: { approvalId: resolution.request.approvalId, toolCallId: resolution.request.toolCallId, status: resolution.decision.status, resolvedAt: resolution.decision.resolvedAt } });
    }, (request) => {
      this.events.append(request.sessionId, { type: "approval.required", runId: request.runId, payload: request });
    });
  }
  setFactory(factory: AgentSessionFactory): void { this.factory = factory; }

  async create(input: CreateSessionInput = {}): Promise<HarnessSession> {
    const sessionId = input.id ?? `ses_${randomUUID()}`;
    const cwd = path.join(this.sessionsRoot, sessionId);
    await mkdir(cwd, { recursive: false, mode: 0o700 });
    const workspace = await SafeWorkspace.create(cwd);
    const session: InternalSession = { sessionId, cwd: workspace.root, workspace, busy: false, currentRun: undefined, agent: undefined, unsubscribe: undefined, timeout: undefined, disposed: false, requestedToolCalls: new Set() };
    this.sessions.set(sessionId, session);
    this.events.append(sessionId, { type: "session.started", payload: { cwd: `workspace/sessions/${sessionId}` } });
    return session;
  }
  get(id: string): HarnessSession | undefined { return this.sessions.get(id); }
  private internal(id: string): InternalSession { const session = this.sessions.get(id); if (!session) throw new HarnessError("SESSION_NOT_FOUND", "Session not found", 404); return session; }

  async configureModel(id: string, input: ModelConnectionInput) {
    const session = this.internal(id);
    if (session.busy) throw new HarnessError("SESSION_BUSY", "Cannot change model configuration while a run is active", 409);
    await this.disposeAgent(session);
    const value = await this.models.configure(id, input);
    this.events.append(id, { type: "model.configured", payload: { providerId: input.providerId, modelId: input.modelId, api: input.api, supportsImages: input.supportsImages } });
    return value;
  }
  async clearCredentials(id: string): Promise<void> { const session = this.internal(id); if (session.busy) throw new HarnessError("SESSION_BUSY", "Cannot clear credentials while a run is active", 409); await this.disposeAgent(session); await this.models.clearCredentials(id); }

  async send(id: string, message: SendMessageInput): Promise<{ runId: string }> {
    const session = this.internal(id);
    if (session.busy) throw new HarnessError("SESSION_BUSY", "A run is already active for this session", 409);
    const configuredModel = this.models.getRedacted(id);
    if (!configuredModel.configured) throw new HarnessError("MODEL_NOT_CONFIGURED", "No model connection is configured", 422);
    if (!configuredModel.keyConfigured) throw new HarnessError("MODEL_CREDENTIAL_NOT_CONFIGURED", "The selected provider has no API key", 422);
    const maxImages = Number(process.env.MAX_IMAGES_PER_MESSAGE ?? 4);
    if (message.attachmentIds.length > maxImages) throw new HarnessError("TOO_MANY_ATTACHMENTS", `A message can include at most ${maxImages} images`, 413);
    const selected = message.attachmentIds.map((attachmentId) => {
      const item = this.attachments.getStored(id, attachmentId);
      if (!item) throw new HarnessError("ATTACHMENT_NOT_FOUND", "Attachment not found for this session", 404);
      return item;
    });
    const expanded = await this.catalog.expandCommand(message.text);
    const runId = `run_${randomUUID()}`;
    const assistantMessageId = `msg_${randomUUID()}`;
    const controller = new AbortController();
    session.currentRun = { runId, controller, assistantMessageId };
    session.busy = true;
    try { await this.ensureAgent(session); } catch (error) { session.busy = false; session.currentRun = undefined; this.events.append(id, { type: "run.failed", runId, payload: { code: "AGENT_SESSION_CREATE_FAILED", message: safeErrorMessage(error) } }); throw error; }
    if (selected.length > 0 && !session.agent?.supportsImages) {
      session.busy = false;
      session.currentRun = undefined;
      throw new HarnessError("MODEL_DOES_NOT_SUPPORT_IMAGES", "The selected model does not support image input", 422);
    }
    const agent = session.agent;
    if (!agent) throw new HarnessError("AGENT_SESSION_CREATE_FAILED", "Agent session was not created", 500);
    this.events.append(id, { type: "user.message", runId, payload: { messageId: `msg_${randomUUID()}`, text: message.text, attachmentIds: message.attachmentIds } });
    this.events.append(id, { type: "run.started", runId, payload: { modelId: agent.modelId } });
    const images: ImageContent[] = await Promise.all(selected.map(async (item) => ({ type: "image" as const, mimeType: item.mediaType, data: await this.attachments.readBase64(id, item.attachmentId) })));
    session.timeout = setTimeout(() => void this.cancel(id, "timeout"), this.runTimeoutMs);
    if (expanded.skill) {
      for (const type of ["skill.requested", "skill.loaded"] as const) this.events.append(id, { type, runId, payload: { name: expanded.skill.name, filePath: expanded.skill.filePath, source: "command" } });
    }
    void this.executeRun(session, redactServerSecret(expanded.text), images);
    return { runId };
  }

  private async ensureAgent(session: InternalSession): Promise<void> {
    if (session.agent) return;
    if (!this.factory) throw new Error("Agent session factory is not configured");
    session.agent = await this.factory.create({
      sessionId: session.sessionId,
      cwd: session.cwd,
      workspace: session.workspace,
      getRun: () => session.currentRun ? { runId: session.currentRun.runId, signal: session.currentRun.controller.signal, assistantMessageId: session.currentRun.assistantMessageId } : undefined,
      skillRequested: (skill, toolCallId) => { const run = session.currentRun; if (run) this.events.append(session.sessionId, { type: "skill.requested", runId: run.runId, payload: { name: skill.name, filePath: skill.filePath, toolCallId, source: "read" } }); },
      skillLoaded: (skill, toolCallId) => { const run = session.currentRun; if (run) this.events.append(session.sessionId, { type: "skill.loaded", runId: run.runId, payload: { name: skill.name, filePath: skill.filePath, toolCallId, source: "read" } }); },
      toolRequested: (toolCallId, toolName, argumentsValue) => { const run = session.currentRun; if (run && !session.requestedToolCalls.has(toolCallId)) { session.requestedToolCalls.add(toolCallId); this.events.append(session.sessionId, { type: "tool.requested", runId: run.runId, payload: { toolCallId, toolName, arguments: redactValue(argumentsValue) } }); } },
    });
    session.unsubscribe = session.agent.subscribe((event) => this.onPiEvent(session, event as AgentSessionEvent));
  }
  private onPiEvent(session: InternalSession, event: AgentSessionEvent): void {
    const run = session.currentRun;
    if (!run) return;
    // The blocking Extension preflight is the single source for tool.requested.
    // Pi's later execution-start event has the same toolCallId and would create
    // duplicate assistant-ui Tool parts (and can expose pre-normalized paths).
    if (event.type === "tool_execution_start") return;
    for (const mapped of mapPiEvent(event, { runId: run.runId, assistantMessageId: run.assistantMessageId })) this.events.append(session.sessionId, mapped);
  }
  private async executeRun(session: InternalSession, text: string, images: ImageContent[]): Promise<void> {
    const run = session.currentRun;
    if (!run || !session.agent) return;
    try {
      await session.agent.prompt(text, images.length ? { images } : undefined);
      if (!run.controller.signal.aborted) this.events.append(session.sessionId, { type: "run.completed", runId: run.runId, payload: {} });
    } catch (error) {
      if (!run.controller.signal.aborted) this.events.append(session.sessionId, { type: "run.failed", runId: run.runId, payload: { code: "AGENT_RUN_FAILED", message: safeErrorMessage(error) } });
    } finally {
      if (session.timeout) clearTimeout(session.timeout);
      session.timeout = undefined;
      session.busy = false;
      if (session.currentRun?.runId === run.runId) session.currentRun = undefined;
    }
  }
  async cancel(id: string, reason: "user" | "timeout" | "session_deleted" | "shutdown" = "user"): Promise<void> {
    const session = this.internal(id);
    const run = session.currentRun;
    if (!run || !session.busy || run.controller.signal.aborted) return;
    run.cancelReason = reason;
    run.controller.abort(reason);
    this.approvals.cancelForSession(id);
    await session.agent?.abort().catch(() => undefined);
    this.events.append(id, { type: "run.cancelled", runId: run.runId, payload: { reason } });
  }
  async delete(id: string, endReason: "deleted" | "shutdown" = "deleted"): Promise<void> {
    const session = this.internal(id);
    if (session.busy) await this.cancel(id, endReason === "shutdown" ? "shutdown" : "session_deleted");
    this.approvals.cancelForSession(id);
    this.events.append(id, { type: "session.ended", payload: { reason: endReason } });
    await this.disposeAgent(session);
    await this.attachments.removeAll(id);
    await this.models.remove(id);
    const root = await realpath(this.sessionsRoot);
    const target = await realpath(session.cwd);
    const relative = path.relative(root, target);
    if (!relative || relative.startsWith("..") || path.isAbsolute(relative)) throw new HarnessError("UNSAFE_SESSION_DELETE", "Refusing unsafe session directory deletion", 500);
    await rm(target, { recursive: true, force: true });
    session.disposed = true;
    this.sessions.delete(id);
  }
  async shutdown(): Promise<void> { for (const id of [...this.sessions.keys()]) await this.delete(id, "shutdown"); }
  private async disposeAgent(session: InternalSession): Promise<void> { session.unsubscribe?.(); session.unsubscribe = undefined; session.agent?.dispose(); session.agent = undefined; }
}

function publicToolArguments(value: unknown, cwd: string, skill?: SkillIndexEntry): unknown {
  const redacted = redactValue(value);
  if (!redacted || typeof redacted !== "object" || Array.isArray(redacted)) return redacted;
  const result = { ...(redacted as Record<string, unknown>) };
  for (const key of ["path", "cwd", "filePath"]) {
    const candidate = result[key];
    if (typeof candidate !== "string" || !path.isAbsolute(candidate)) continue;
    if (skill && key === "path") result[key] = skill.filePath;
    else {
      const relative = path.relative(cwd, candidate);
      result[key] = relative && !relative.startsWith("..") && !path.isAbsolute(relative) ? relative : "<protected-path>";
    }
  }
  return result;
}