import type { TimelineItem } from "./streaming-v2";
import type { InputOutcome, WorkItemDiagnostic } from "./protocol-v2";

export type InputItem = Extract<TimelineItem, { lane: "input" }>;
export type CellItem = Extract<TimelineItem, { lane: "cell" }>;
export type SpeechItem = Extract<TimelineItem, { lane: "speech" }>;
export type LifecycleItem = Extract<TimelineItem, { lane: "lifecycle" }>;

/** Consecutive cell executions collapse into one row: twenty identical cells is not information. */
export interface CellRun {
    kind: "cells";
    id: string;
    items: CellItem[];
}

/**
 * The same rule for the bookkeeping lane. `model usage` repeated six times as six separate lines is
 * six rows of the same sentence; one row that says how many, and which kinds, is the same fact.
 */
export interface LifecycleRun {
    kind: "lifecycle";
    id: string;
    items: LifecycleItem[];
}

export type TurnChild = TimelineItem | CellRun | LifecycleRun;

export function isCellRun(child: TurnChild): child is CellRun {
    return (child as CellRun).kind === "cells";
}

export function isLifecycleRun(child: TurnChild): child is LifecycleRun {
    return (child as LifecycleRun).kind === "lifecycle";
}

export interface TurnSummary {
    cells: number;
    ok: number;
    failed: number;
    delivered: number;
    held: number;
    rejected: number;
    scratch: number;
    lifecycle: number;
}

export interface Turn {
    id: string;
    /** The human input that opened this turn, when there was one. */
    input?: InputItem;
    startedAt: number;
    /**
     * Where this turn sits in the session's own log: the sequence of its opening record. A console that
     * loads a window cannot count turns from one — that would number turn 361 as "#1" and tell the
     * reader the session began where their window does — so a turn is named by the coordinate that is
     * the same whatever part of the session is loaded.
     */
    seq?: number;
    children: TurnChild[];
    summary: TurnSummary;
}

/**
 * Which lanes the operator is looking at. A set, not a single choice: "conversation" is a named
 * preset rather than a mode, so you can watch your own words next to the machine's output without
 * being forced into one view.
 */
export type LaneSelection = string[];

export const LANE_IDS = ["speech", "cell", "scratchpad", "lifecycle", "input"] as const;

export function isSelected(selection: LaneSelection, id: string): boolean {
    return selection.includes(id);
}

/** Everything, expressed as a selection. */
export function everything(): LaneSelection {
    return ["conversation", ...LANE_IDS];
}

/**
 * The view a reader wants without choosing: what you said, everything the agent said whether or not
 * the delivery review let it through, and what it ran. `conversation` is redundant beside `speech`
 * (it only adds your own inputs) but it is what makes the inputs appear, so it stays.
 */
export const TRANSCRIPT_PRESET: LaneSelection = ["conversation", "speech", "cell"];

export function sameSelection(a: LaneSelection, b: LaneSelection): boolean {
    return a.length === b.length && a.every((id) => b.includes(id));
}

export function filterItems(items: TimelineItem[], selection: LaneSelection): TimelineItem[] {
    if (selection.length === 0) return [];
    const wants = new Set(selection);
    const all = wants.has("conversation") && LANE_IDS.every((id) => wants.has(id));
    if (all) return items;

    return items.filter((item) => {
        if (wants.has(item.lane)) return true;
        // the conversation preset: your own words, plus the speech that was actually delivered
        if (wants.has("conversation")) {
            if (item.lane === "input") return true;
            if (item.lane === "speech" && item.data.receipt?.status === "confirmed") return true;
        }
        return false;
    });
}

/** True when this record belongs to the conversation preset but not to a lane the operator picked. */
export function onlyViaConversation(item: TimelineItem, selection: LaneSelection): boolean {
    return selection.includes("conversation") && !selection.includes(item.lane);
}

function receiptOf(item: TimelineItem): string | undefined {
    return item.lane === "speech" ? item.data.receipt?.status : undefined;
}

function summarise(items: TimelineItem[]): TurnSummary {
    const summary: TurnSummary = { cells: 0, ok: 0, failed: 0, delivered: 0, held: 0, rejected: 0, scratch: 0, lifecycle: 0 };
    for (const item of items) {
        switch (item.lane) {
            case "cell": {
                summary.cells += 1;
                const status = item.data.status;
                if (status === "ok") summary.ok += 1;
                else if (status === "error" || item.data.exit_code) summary.failed += 1;
                break;
            }
            case "speech": {
                const receipt = receiptOf(item);
                if (receipt === "confirmed") summary.delivered += 1;
                else if (receipt === "held") summary.held += 1;
                else if (receipt === "rejected" || receipt === "failed") summary.rejected += 1;
                break;
            }
            case "scratchpad":
                summary.scratch += 1;
                break;
            case "lifecycle":
                summary.lifecycle += 1;
                break;
            default:
                break;
        }
    }
    return summary;
}

/** Collapse runs of consecutive cell executions, keeping every other record in order. */
/** Runs of like records collapse; a run of one is just the record. */
function collapseRuns(items: TimelineItem[]): TurnChild[] {
    const children: TurnChild[] = [];
    let cells: CellItem[] = [];
    let life: LifecycleItem[] = [];

    const flushCells = () => {
        if (cells.length === 0) return;
        if (cells.length === 1) children.push(cells[0]);
        else children.push({ kind: "cells", id: `run-${cells[0].id}`, items: [...cells] });
        cells = [];
    };
    const flushLife = () => {
        if (life.length === 0) return;
        if (life.length === 1) children.push(life[0]);
        else children.push({ kind: "lifecycle", id: `life-${life[0].id}`, items: [...life] });
        life = [];
    };

    for (const item of items) {
        if (item.lane === "cell") {
            flushLife();
            cells.push(item);
            continue;
        }
        if (item.lane === "lifecycle") {
            flushCells();
            life.push(item);
            continue;
        }
        flushCells();
        flushLife();
        children.push(item);
    }
    flushCells();
    flushLife();
    return children;
}

/**
 * What became of each thing the operator said. The answer for a turn is the evidence between it and the
 * next input — the delivered speech, the held speech, the failing cell — and one exact answer that
 * outranks all of it: a `turn.outcome` record naming the input's own event ids.
 *
 * The evidence is gathered in the single walk that finds the segments, not by scanning the transcript
 * once per input. It used to scan per input, and this runs on every timeline change: with a few thousand
 * records that scan was the largest single cost in the console, larger than the whole render.
 */
export function computeInputOutcomes(
    items: TimelineItem[],
    work: WorkItemDiagnostic[] = []
): Map<string, InputOutcome> {
    // First pass: the answers that are stated outright, keyed by the input sequence they name.
    const exactBySeq = new Map<number, InputOutcome>();
    for (const item of items) {
        if (item.lane !== "lifecycle") continue;
        const kind = (item.data as { kind?: string }).kind;
        if (kind !== "turn.outcome" && kind !== "turn.failed") continue;
        const data = (item.data as { data?: Record<string, unknown> }).data ?? {};
        const ids = Array.isArray(data.input_event_ids) ? (data.input_event_ids as number[]) : [];
        const status = String(data.status ?? (kind === "turn.failed" ? "failed" : ""));
        const reason = String(data.error ?? data.reason ?? (item.data as { reason?: string }).reason ?? "");
        const outcome: InputOutcome =
            kind === "turn.failed" || status === "failed"
                ? { status: "failed", reason: reason || "the turn failed" }
                : status === "held"
                ? { status: "held", reason: reason || "held by delivery review" }
                : status === "replied"
                ? { status: "replied" }
                : { status: "no_reply", reason: "the turn completed without speech" };
        for (const id of ids) exactBySeq.set(id, outcome);
    }

    // Second pass: walk the transcript, keeping the segment started by the open input. Each input is
    // decided when the next one closes it, so every record is read exactly once.
    const failedWork = work.filter((item) => item.status === "failed");
    const outcomes = new Map<string, InputOutcome>();
    let open: InputItem | null = null;
    let delivered = false;
    let heldReason: string | null = null;
    let failedCell: CellItem | null = null;
    let failedLifecycle = false;

    const close = (last: boolean) => {
        const input = open;
        if (!input) return;
        open = null;
        if (input.seq !== undefined && exactBySeq.has(input.seq)) {
            outcomes.set(input.id, exactBySeq.get(input.seq)!);
        } else if (delivered) {
            outcomes.set(input.id, { status: "replied" });
        } else if (heldReason !== null) {
            outcomes.set(input.id, { status: "held", reason: heldReason });
        } else {
            const failed = failedCell
                ? failedCell.data.outputs.find((output) => output.stream === "stderr")?.text
                    ?.split("\n")
                    .filter(Boolean)
                    .slice(-1)[0] || `the cell exited ${failedCell.data.exit_code ?? "non-zero"}`
                : failedWork.find((item) => input.seq !== undefined && item.id.includes(`input-${input.seq}`))?.title;
            if (failed || failedLifecycle) {
                outcomes.set(input.id, { status: "failed", reason: failed || "the turn failed" });
            } else if (last) {
                outcomes.set(input.id, { status: "running" });
            } else {
                outcomes.set(input.id, { status: "no_reply", reason: "the turn completed without speech" });
            }
        }
        delivered = false;
        heldReason = null;
        failedCell = null;
        failedLifecycle = false;
    };

    for (const item of items) {
        if (item.lane === "input") {
            close(false);
            open = item;
            continue;
        }
        if (!open) continue;
        switch (item.lane) {
            case "speech": {
                const receipt = item.data.receipt;
                if (receipt?.status === "confirmed") delivered = true;
                else if (receipt?.status === "held" && heldReason === null) {
                    heldReason = receipt.reason || "held by delivery review";
                }
                break;
            }
            case "cell":
                if (failedCell === null && (item.data.status === "error" || (item.data.exit_code ?? 0) !== 0)) {
                    failedCell = item;
                }
                break;
            case "lifecycle":
                if (String((item.data as { kind?: string }).kind) === "failed") failedLifecycle = true;
                break;
            default:
                break;
        }
    }
    close(true);
    return outcomes;
}

/**
 * A turn is one human input and everything the runtime did about it. Records that arrive before any
 * input (session activation, a model-driven start) form a leading turn of their own, so nothing is
 * dropped and nothing is invented.
 */
export function groupTurns(items: TimelineItem[]): Turn[] {
    const turns: Turn[] = [];
    let current: TimelineItem[] = [];
    let input: InputItem | undefined;

    const close = () => {
        if (current.length === 0 && !input) return;
        turns.push({
            id: input?.id ?? current[0]?.id ?? `turn-${turns.length}`,
            input,
            startedAt: input?.created_ms ?? current[0]?.created_ms ?? Date.now(),
            seq: input?.seq ?? current[0]?.seq,
            children: collapseRuns(current),
            summary: summarise(current),
        });
        current = [];
        input = undefined;
    };

    for (const item of items) {
        if (item.lane === "input") {
            close();
            input = item;
            continue;
        }
        current.push(item);
    }
    close();
    return turns;
}

/**
 * One row is one line of the transcript. The transcript is a tree — turn, then the records inside it,
 * then the cells inside an open run — and a virtualizer can only place a flat list. The tree is
 * flattened here, in the data layer, so the view receives rows that already know their own key, their
 * own spacing and their own timestamp. Doing it here also makes it testable without a browser.
 */
export type Row = {
    key: string;
    turnId: string;
    /** The moment this row reports to the horizon readout. */
    at: number;
    /** The last row of its turn: it carries the turn's own bottom gap. */
    last: boolean;
    /** A guess until the row has been measured. Only ever used for rows off screen. */
    estimate: number;
} & (
    /**
     * The top of the loaded window. A transcript is read from its live end, so the records above it are
     * on the daemon, not here: this row is what asks for them, and it is a row rather than a floating
     * control so it can never be mistaken for a record of the session.
     */
    | { kind: "history" }
    /**
     * `partial` marks the one turn a window can cut in half: the first, when the records that opened it
     * are above the window. Its counts are counts of the records that are loaded, which is not the same
     * claim as a turn's summary, so it makes none.
     */
    | { kind: "head"; turn: Turn; no: number; partial?: boolean }
    | { kind: "input"; turn: Turn; item: InputItem }
    | { kind: "child"; turn: Turn; child: TurnChild }
    | { kind: "cell"; turn: Turn; cell: CellItem; endOfRun: boolean }
);

const ESTIMATE = { head: 44, input: 56, cell: 96, record: 52 } as const;

/** A collapsed run is one line; a lone cell is a terminal; everything else is a short record. */
function estimateOf(child: TurnChild): number {
    if (isCellRun(child) || isLifecycleRun(child)) return ESTIMATE.record;
    return child.lane === "cell" ? ESTIMATE.cell : ESTIMATE.record;
}

function atOf(child: TurnChild): number {
    if (isCellRun(child) || isLifecycleRun(child)) return child.items[0]?.created_ms ?? 0;
    return child.created_ms;
}

/**
 * `open` is the set of expanded cell runs. An expanded run emits one row per cell, which is the whole
 * point of building rows here: a turn that ran two hundred cells is two hundred rows the scroller can
 * drop, not one enormous element it must keep.
 */
export function buildRows(turns: Turn[], open: Record<string, boolean> = {}, earlier = false): Row[] {
    const rows: Row[] = [];
    // Only the top of the loaded window can cut a turn: every page below it joins exactly where the
    // page above it ended, so a turn split by one page is whole again once the next one arrives.
    const cut = earlier && turns.length > 0 && !turns[0].input;
    if (earlier) {
        // `at: 0` keeps it out of the horizon readout, which reports the moment the reader is looking at.
        rows.push({ kind: "history", key: "history", turnId: "", at: 0, last: false, estimate: 40 });
    }
    for (const [i, turn] of turns.entries()) {
        // A record with no sequence is synthetic (the tests build them): rank is all it has.
        rows.push({
            kind: "head", key: `${turn.id}:head`, turnId: turn.id, at: turn.startedAt, last: false,
            estimate: ESTIMATE.head, turn, no: turn.seq ?? i + 1, partial: cut && i === 0,
        });
        if (turn.input) {
            rows.push({ kind: "input", key: turn.input.id, turnId: turn.id, at: turn.input.created_ms, last: false, estimate: ESTIMATE.input, turn, item: turn.input });
        }
        for (const child of turn.children) {
            const at = atOf(child);
            if (isCellRun(child) && open[child.id]) {
                rows.push({ kind: "child", key: `${child.id}:head`, turnId: turn.id, at, last: false, estimate: ESTIMATE.record, turn, child });
                for (const [j, cell] of child.items.entries()) {
                    rows.push({ kind: "cell", key: cell.id, turnId: turn.id, at: cell.created_ms, last: false, estimate: ESTIMATE.cell, turn, cell, endOfRun: j === child.items.length - 1 });
                }
                continue;
            }
            rows.push({ kind: "child", key: child.id, turnId: turn.id, at, last: false, estimate: estimateOf(child), turn, child });
        }
    }
    // spacing belongs to the boundary between turns, so both rows that touch it need to know
    for (let i = 1; i < rows.length; i++) {
        if (rows[i].turnId !== rows[i - 1].turnId) rows[i - 1].last = true;
    }
    if (rows.length > 0) rows[rows.length - 1].last = true;
    return rows;
}