import type { ChatEvent, TurnMessage } from '../../api/chat/types';
import { appendChatEvent, getLastChatEventSeq } from 'officerdb';
import { createTurnStream, type TurnOutput } from './turn-stream';

// The agent sidecar is the writer of record for chat output.
//
// It used to push raw ChatEvents at officer over the registration socket and let officer translate and
// persist them. That socket silently drops when officer is down (`connect.ts:send` — no queue, no
// error), so everything the agent produced during a restart was lost: the turn kept running here and
// its output went nowhere. Writing to Postgres here instead means an officer restart costs a replay
// rather than the output, because the durable record no longer travels over the socket that died.
//
// Officer still gets every message live — it just gets it already written, with its cursor id attached,
// and relays it verbatim.

export type Delivery = {
  sessionId: string;
  msg: TurnMessage;
  /** The cursor id under which this message is durable. Absent = ephemeral delta, or the write failed. */
  seq?: number;
};

export type SessionLogStore = {
  /**
   * Translate one parser event and commit + deliver whatever it produces, strictly in order.
   * `durable: false` skips the write (see ClaudeSpawnStreamingParams.durable) — the messages are still
   * delivered, just without a cursor, because nothing will ever replay them.
   */
  push: (sessionId: string, event: ChatEvent, durable?: boolean) => void;
  /** Forget a session's buffer and cursor chain (on kill / clear-session). */
  drop: (sessionId: string) => void;
  /**
   * Settle every commit queued so far. For shutdown: `push` returns immediately and the durable write
   * happens on a queue, so exiting without this drops whatever had not reached Postgres yet.
   */
  flush: () => Promise<void>;
};

// The durable store, behind an interface so the ordering guarantee below can be tested against a writer
// whose writes finish out of order. Defaults to Postgres.
export type EventWriter = {
  append: (sessionId: string, msg: TurnMessage) => Promise<number>;
  lastSeq: (sessionId: string) => Promise<number | undefined>;
};

const postgresWriter: EventWriter = {
  append: (sessionId, msg) => appendChatEvent(sessionId, msg),
  lastSeq: (sessionId) => getLastChatEventSeq(sessionId),
};

type SessionLog = {
  stream: ReturnType<typeof createTurnStream>;
  /** Serialises commits so cursor ids are assigned in the order the events actually arrived. */
  tail: Promise<void>;
  lastSeq: number | undefined;
  resolvedLastSeq: boolean;
};

export function createSessionLogStore(
  deliver: (d: Delivery) => void,
  writer: EventWriter = postgresWriter,
): SessionLogStore {
  const logs = new Map<string, SessionLog>();

  function logFor(sessionId: string): SessionLog {
    let log = logs.get(sessionId);
    if (!log) {
      log = {
        stream: createTurnStream(sessionId),
        tail: Promise.resolve(),
        lastSeq: undefined,
        resolvedLastSeq: false,
      };
      logs.set(sessionId, log);
    }
    return log;
  }

  async function commit(sessionId: string, log: SessionLog, out: TurnOutput, durable: boolean): Promise<void> {
    // Deltas are live-only, but still go through the queue: a delta that overtook the `assistant:text`
    // or `tool:start` in front of it would make the client commit its stream buffer at the wrong point.
    if (!out.durable || !durable) {
      deliver({ sessionId, msg: out.msg });
      return;
    }

    // Pick the chain back up after a restart of this process, so `prevSeq` stays meaningful for a
    // session that started before it. Once per session; failure just means no continuity claim.
    if (!log.resolvedLastSeq) {
      log.resolvedLastSeq = true;
      try {
        log.lastSeq = await writer.lastSeq(sessionId);
      } catch (err) {
        console.error(`[agent] could not read last event cursor for ${sessionId}:`, err);
      }
    }

    const msg: TurnMessage = log.lastSeq === undefined ? out.msg : { ...out.msg, prevSeq: log.lastSeq };

    try {
      const seq = await writer.append(sessionId, msg);
      log.lastSeq = seq;
      deliver({ sessionId, msg, seq });
    } catch (err) {
      // The write IS the durability guarantee, so don't pretend. Deliver live without a cursor: the
      // client sees the message but won't advance its cursor past something it cannot replay, and the
      // next successful write chains from the last cursor the client actually holds.
      console.error(`[agent] failed to persist chat event for ${sessionId}:`, err);
      deliver({ sessionId, msg: out.msg });
    }
  }

  return {
    push(sessionId, event, durable = true) {
      const log = logFor(sessionId);
      // Translation is synchronous and therefore in arrival order; only the commit is queued.
      for (const out of log.stream.push(event)) {
        log.tail = log.tail.then(() => commit(sessionId, log, out, durable));
      }
    },

    drop(sessionId) {
      logs.delete(sessionId);
    },

    async flush() {
      // Await the tails as they stand now. A `push` racing this one is not waited for — which is the
      // honest contract: this exists for shutdown, where the alternative is exiting on an unresolved
      // write and losing the row that explains why the turn ended.
      await Promise.allSettled([...logs.values()].map((log) => log.tail));
    },
  };
}