import { EventEmitter } from "node:events";
import crypto from "node:crypto";
import type { WorkflowId } from "./workflows.js";

// ── Types ──────────────────────────────────────────────────────────

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

export interface Message {
  id: number;
  sessionId: string;
  from: AgentId;
  to: AgentId;
  body: string;
  status: "pending" | "delivered";
  createdAt: number;
}

export interface Session {
  id: string;
  goal: string;
  initiator: AgentId;
  responder: AgentId;
  workflow: WorkflowId;
  currentPhaseIndex: number;
  status: "active" | "closed";
  autoForward: boolean;
  lastSeenByAgent: Record<AgentId, number>;
  messages: Message[];
  nextMessageId: number;
  createdAt: number;
  closedAt?: number;
  closedBy?: AgentId | "user";
  closedReason?: string;
}

export type UIEvent =
  | { type: "session-created"; session: Session }
  | { type: "session-updated"; session: Session }
  | { type: "session-deleted"; sessionId: string }
  | { type: "message-added"; sessionId: string; message: Message }
  | { type: "message-updated"; sessionId: string; message: Message }
  | { type: "message-removed"; sessionId: string; messageId: number };

export type WaitForMessageResult =
  | { status: "received"; message: Message; skippedMessageIds: number[] }
  | { status: "timeout" }
  | { status: "closed" };

// ── Store ──────────────────────────────────────────────────────────

export class Store extends EventEmitter {
  private sessions = new Map<string, Session>();

  createSession(goal: string, initiator: AgentId, workflow: WorkflowId): Session {
    const id = "collab-" + crypto.randomBytes(4).toString("hex");
    const responder: AgentId = initiator === "claude" ? "codex" : "claude";
    const session: Session = {
      id,
      goal,
      initiator,
      responder,
      workflow,
      currentPhaseIndex: 0,
      status: "active",
      autoForward: true,
      lastSeenByAgent: { claude: 0, codex: 0 },
      messages: [],
      nextMessageId: 1,
      createdAt: Date.now(),
    };
    this.sessions.set(id, session);
    this.emit("ui-update", { type: "session-created", session } satisfies UIEvent);
    return session;
  }

  getSession(id: string): Session | undefined {
    return this.sessions.get(id);
  }

  listSessions(): Session[] {
    return [...this.sessions.values()];
  }

  updateSession(
    id: string,
    updates: { autoForward?: boolean; goal?: string; currentPhaseIndex?: number },
  ): Session | undefined {
    const session = this.sessions.get(id);
    if (!session) return undefined;
    if (updates.autoForward !== undefined) session.autoForward = updates.autoForward;
    if (updates.goal !== undefined) session.goal = updates.goal;
    if (updates.currentPhaseIndex !== undefined) session.currentPhaseIndex = updates.currentPhaseIndex;
    this.emit("ui-update", { type: "session-updated", session } satisfies UIEvent);
    return session;
  }

  deleteSession(id: string): boolean {
    const existed = this.sessions.delete(id);
    if (existed) {
      this.emit("ui-update", { type: "session-deleted", sessionId: id } satisfies UIEvent);
    }
    return existed;
  }

  closeSession(sessionId: string, closedBy: AgentId | "user", reason?: string): Session | undefined {
    const session = this.sessions.get(sessionId);
    if (!session) return undefined;
    if (session.status === "closed") return session;

    session.status = "closed";
    session.closedAt = Date.now();
    session.closedBy = closedBy;
    if (reason && reason.trim()) session.closedReason = reason.trim();

    this.emit("ui-update", { type: "session-updated", session } satisfies UIEvent);
    this.emit(`message:${sessionId}`);
    return session;
  }

  addMessage(sessionId: string, from: AgentId, body: string): Message | undefined {
    const session = this.sessions.get(sessionId);
    if (!session) return undefined;
    if (session.status === "closed") return undefined;

    const to: AgentId = from === "claude" ? "codex" : "claude";
    const msg: Message = {
      id: session.nextMessageId++,
      sessionId,
      from,
      to,
      body,
      status: session.autoForward ? "delivered" : "pending",
      createdAt: Date.now(),
    };
    session.messages.push(msg);

    this.emit("ui-update", { type: "message-added", sessionId, message: msg } satisfies UIEvent);

    if (msg.status === "delivered") {
      this.emit(`message:${sessionId}`, msg);
    }

    return msg;
  }

  getLatestUnseenOutbound(sessionId: string, from: AgentId): Message | undefined {
    const session = this.sessions.get(sessionId);
    if (!session) return undefined;

    const to: AgentId = from === "claude" ? "codex" : "claude";
    const seenId = session.lastSeenByAgent[to] ?? 0;
    for (let i = session.messages.length - 1; i >= 0; i--) {
      const m = session.messages[i];
      if (m.from === from && m.to === to && m.id > seenId) {
        return m;
      }
    }
    return undefined;
  }

  deliverMessage(sessionId: string, messageId: number): Message | undefined {
    const session = this.sessions.get(sessionId);
    if (!session) return undefined;

    const msg = session.messages.find((m) => m.id === messageId);
    if (!msg || msg.status !== "pending") return undefined;

    msg.status = "delivered";
    this.emit("ui-update", { type: "message-updated", sessionId, message: msg } satisfies UIEvent);
    this.emit(`message:${sessionId}`, msg);
    return msg;
  }

  editMessage(sessionId: string, messageId: number, newBody: string): Message | undefined {
    const session = this.sessions.get(sessionId);
    if (!session) return undefined;

    const msg = session.messages.find((m) => m.id === messageId);
    if (!msg || msg.status !== "pending") return undefined;

    msg.body = newBody;
    this.emit("ui-update", { type: "message-updated", sessionId, message: msg } satisfies UIEvent);
    return msg;
  }

  rejectMessage(sessionId: string, messageId: number): boolean {
    const session = this.sessions.get(sessionId);
    if (!session) return false;

    const idx = session.messages.findIndex((m) => m.id === messageId && m.status === "pending");
    if (idx === -1) return false;

    session.messages.splice(idx, 1);
    this.emit("ui-update", { type: "message-removed", sessionId, messageId } satisfies UIEvent);
    return true;
  }

  getLatestDeliveredForAgent(sessionId: string, agent: AgentId, afterId: number): Message | undefined {
    const session = this.sessions.get(sessionId);
    if (!session) return undefined;

    for (let i = session.messages.length - 1; i >= 0; i--) {
      const m = session.messages[i];
      if (m.to === agent && m.status === "delivered" && m.id > afterId) {
        return m;
      }
    }
    return undefined;
  }

  getSkippedDeliveredIdsForAgent(sessionId: string, agent: AgentId, afterId: number, latestId: number): number[] {
    const session = this.sessions.get(sessionId);
    if (!session) return [];

    return session.messages
      .filter((m) => m.to === agent && m.status === "delivered" && m.id > afterId && m.id < latestId)
      .map((m) => m.id);
  }

  markSeen(sessionId: string, agent: AgentId, messageId: number): void {
    const session = this.sessions.get(sessionId);
    if (!session) return;
    session.lastSeenByAgent[agent] = Math.max(session.lastSeenByAgent[agent] ?? 0, messageId);
  }

  waitForMessage(
    sessionId: string,
    agent: AgentId,
    afterId: number,
    timeoutMs: number,
  ): Promise<WaitForMessageResult> {
    const buildReceived = (msg: Message): WaitForMessageResult => {
      const skippedMessageIds = this.getSkippedDeliveredIdsForAgent(sessionId, agent, afterId, msg.id);
      this.markSeen(sessionId, agent, msg.id);
      return { status: "received", message: msg, skippedMessageIds };
    };

    const session = this.sessions.get(sessionId);
    if (session?.status === "closed") return Promise.resolve({ status: "closed" });

    const existing = this.getLatestDeliveredForAgent(sessionId, agent, afterId);
    if (existing) return Promise.resolve(buildReceived(existing));

    return new Promise((resolve) => {
      let settled = false;
      const eventName = `message:${sessionId}`;

      const onMessage = () => {
        const current = this.sessions.get(sessionId);
        if (current?.status === "closed" && !settled) {
          settled = true;
          clearTimeout(timer);
          this.removeListener(eventName, onMessage);
          resolve({ status: "closed" });
          return;
        }

        const msg = this.getLatestDeliveredForAgent(sessionId, agent, afterId);
        if (msg && !settled) {
          settled = true;
          clearTimeout(timer);
          this.removeListener(eventName, onMessage);
          resolve(buildReceived(msg));
        }
      };

      const timer = setTimeout(() => {
        if (!settled) {
          settled = true;
          this.removeListener(eventName, onMessage);
          resolve({ status: "timeout" });
        }
      }, timeoutMs);

      // Register listener BEFORE re-checking to close the race window
      this.on(eventName, onMessage);

      // Re-check: a message may have arrived between the fast-path check and
      // the listener registration above.
      const arrived = this.getLatestDeliveredForAgent(sessionId, agent, afterId);
      if (arrived && !settled) {
        settled = true;
        clearTimeout(timer);
        this.removeListener(eventName, onMessage);
        resolve(buildReceived(arrived));
      }
    });
  }
}