import { create } from "zustand";
import type { StreamEvent, WidgetDescriptor } from "../api/types";

interface Entry {
  role: "assistant";
  text: string;
}

interface SessionState {
  pendingWidget: WidgetDescriptor | null;
  transcript: Entry[];
  activity: Entry[];
  toasts: { level: string; text: string }[];
  stepMessages: { level: string; text: string }[];
  lastUserEntry: { kind: "input" | "choice"; text: string } | null;
  lastSystemEvent: StreamEvent | null;
  currentPhase: string | null;
  phaseError: string | null;
  // True while a Pi turn is in flight. Turned off by the backend-forwarded `agent_end`
  // system event — the only end-of-turn signal on the final workflow step, where no
  // follow-up gate arrives to release the spinner.
  agentActive: boolean;
  applyEvent: (e: StreamEvent) => void;
  clearPending: () => void;
  resetSession: () => void;
  setPhase: (phase: string | null) => void;
  pushToast: (toast: { level: string; text: string }) => void;
  setLastUserEntry: (e: { kind: "input" | "choice"; text: string }) => void;
  setAgentActive: (v: boolean) => void;
}

const empty = {
  pendingWidget: null,
  transcript: [] as Entry[],
  activity: [] as Entry[],
  toasts: [] as { level: string; text: string }[],
  stepMessages: [] as { level: string; text: string }[],
  lastUserEntry: null as { kind: "input" | "choice"; text: string } | null,
  lastSystemEvent: null,
  currentPhase: null as string | null,
  phaseError: null as string | null,
  agentActive: false,
};

export const useSessionStore = create<SessionState>((set) => ({
  ...empty,
  applyEvent: (e) =>
    set((st) => {
      if (e.type === "ui_request")
        return {
          pendingWidget: e.ui_request,
          // The turn is blocked on the gate, not ended: keep the agent marked active.
          agentActive: true,
          currentPhase: e.ui_request.phase
            ? e.ui_request.phase.split("_")[0]
            : st.currentPhase,
          // A new gate means the phase moved on (or re-presented): clear any error flag.
          phaseError: null,
        };
      if (e.type === "text_delta") {
        const t = [...st.transcript];
        const last = t.at(-1);
        if (last && !st.pendingWidget) t[t.length - 1] = { role: "assistant", text: last.text + e.text };
        else t.push({ role: "assistant", text: e.text });
        // Streamed text means the turn is alive (also covers reattaching mid-turn).
        return { transcript: t, agentActive: true };
      }
      if (e.type === "activity_delta") {
        const activity = [...st.activity];
        const last = activity.at(-1);
        if (last) activity[activity.length - 1] = { role: "assistant", text: last.text + e.text };
        else activity.push({ role: "assistant", text: e.text });
        return { activity, agentActive: true };
      }
      if (e.type === "info") {
        const stepMessages = [...st.stepMessages, { level: e.level ?? "info", text: e.text }];
        // An error during the active phase marks that phase red until the next gate.
        return e.level === "error" ? { stepMessages, phaseError: st.currentPhase } : { stepMessages };
      }
      if (e.type === "system_event")
        return e.event === "agent_end"
          ? { lastSystemEvent: e, agentActive: false }
          : { lastSystemEvent: e };
      return {};
    }),
  clearPending: () => set({ pendingWidget: null }),
  resetSession: () => set({ ...empty }),
  setPhase: (phase) => set({ currentPhase: phase }),
  pushToast: (toast) => set((st) => ({ toasts: [...st.toasts, toast] })),
  // A user entry (question, gate choice, steer) hands the ball back to the harness.
  setLastUserEntry: (e) => set({ lastUserEntry: e, stepMessages: [], agentActive: true }),
  setAgentActive: (v) => set({ agentActive: v }),
}));