import { useCallback, useMemo, useRef } from "react";
import { Effect } from "effect";
import {
  finalizeRunningToolBlocks,
  mergeCanonicalAndRuntimeEvents,
  reconcileReplayMessages,
  replayCursorAfterRuntimeHydration,
  runtimeStatusAcceptsControl,
} from "@/features/agent/messages";
import { foldSessionEvents } from "@/features/agent/runtime/pi-event-applier";
import { settleTurnFinalizingTools } from "@/features/agent/runtime/session-status";
import {
  selectedContextPrompt,
  type ComposerPromptTemplateRef,
  type ComposerSkillRef,
} from "@/features/agent/composer-context";
import type { Session, SessionId, UpdateSession } from "@/features/agent/runtime/types";
import type { BrowserBackend, ToolSelection } from "@/features/agent/tools/types";
import type {
  AgentQueueAction,
  AgentThinkingLevel,
  AgentToolAccess,
} from "@/features/agent/contracts";
import * as api from "@/features/agent/runtime/api";
import {
  runtimeCanHydrateCanonicalSession,
  submitPromptTurn,
  type SubmitArgs,
} from "@/features/agent/runtime/prompt-stream";
import { readTranscriptSnapshot } from "@/features/agent/workspace/transcript-cache";

import { sessionRuntimeController } from "@/features/agent/runtime/session-runtime-controller";

const EMPTY_SKILLS: ComposerSkillRef[] = [];
const EMPTY_PROMPT_TEMPLATES: ComposerPromptTemplateRef[] = [];
const inFlightReplays = new Set<SessionId>();

export type UseSessionEngineDeps = {
  /** Latest `tabs` snapshot — engine reads via a ref so it doesn't restart on every frame. */
  tabs: Session[];
  activeTabId: SessionId;
  modelId: string;
  thinkingLevel: AgentThinkingLevel;
  toolAccess: AgentToolAccess;
  cwd: string;
  browserToolEnabled: boolean;
  browserBackend: BrowserBackend;
  onPiSessionIdChange?: (piSessionId: string) => void;
  /** Mutate a single session record. */
  updateSession: UpdateSession;
  /** Look up the per-session tool selection from the tools subsystem. */
  selectionFor: (sessionId: SessionId) => ToolSelection;
};

export type SessionEngine = {
  /** Send a freshly-typed prompt — orchestrates optimistic update + streaming. */
  submitPrompt: (args: SubmitArgs) => Promise<void>;
  /** Send a steer/follow-up control message while a turn is in progress. */
  sendControl: (request: AgentControlRequest) => Promise<{ ok: boolean; error?: string }>;
  loadRuntimeStatus: (
    runtime: string,
    piSessionId?: string | null,
  ) => Promise<api.RuntimeStatus | null>;
  abortTurn: (sessionId: SessionId) => Promise<api.AbortSessionResult>;
  loadAndReplay: (piSessionId: string, sessionId: SessionId) => Promise<void>;
  /** Fetch and prepend the previous page of older history (tail paging). */
  loadEarlier: (sessionId: SessionId) => Promise<void>;
  compact: (sessionId: SessionId) => Promise<void>;
  /** Probe whether the session's live runtime accepts steer/follow-up right
   * now: running/starting locally, and the runtime's reported pi session (if
   * any) matches ours. A failed probe counts as accepting — the turn API
   * itself is the authority and will reject if not. */
  acceptsControl: (
    tab: { status: Session["status"]; piSessionId?: string | null },
    runtime: string,
  ) => Promise<boolean>;
};

export type AgentControlRequest = {
  mode: "steer" | "follow_up";
  text: string;
  runtime: string;
  sessionId: SessionId;
  piSessionId?: string | null;
  queueAction?: AgentQueueAction;
  queueReplacement?: string;
};

export function useSessionEngine(deps: UseSessionEngineDeps): SessionEngine {
  const {
    tabs,
    activeTabId,
    modelId,
    thinkingLevel,
    toolAccess,
    cwd,
    browserToolEnabled,
    browserBackend,
    onPiSessionIdChange,
    updateSession,
    selectionFor,
  } = deps;

  const tabsRef = useRef(tabs);
  tabsRef.current = tabs;
  const selectionForRef = useRef(selectionFor);
  selectionForRef.current = selectionFor;
  // Sessions with an in-flight "load earlier" page, so a double click / repeated
  // scroll doesn't fetch and prepend the same chunk twice.
  const loadingEarlierRef = useRef<Set<SessionId>>(new Set());

  const loadRuntimeStatusCb = useCallback(api.loadRuntimeStatus, []);

  const sendControl = useCallback(
    (request: AgentControlRequest): Promise<{ ok: boolean; error?: string }> => {
      const { mode, text, runtime, sessionId, piSessionId, queueAction, queueReplacement } =
        request;
      if (!text.trim() || !modelId) return Promise.resolve({ ok: false });
      return Effect.runPromise(
        Effect.gen(function* () {
          const selection = selectionForRef.current(sessionId);
          const skills = selection.skills ?? EMPTY_SKILLS;
          const promptTemplates = selection.promptTemplates ?? EMPTY_PROMPT_TEMPLATES;
          const browserEnabledForTurn = browserToolEnabled;
          const message = selectedContextPrompt(text, skills);
          const contextualQueueReplacement = queueReplacement
            ? selectedContextPrompt(queueReplacement, skills)
            : undefined;
          const result = yield* Effect.tryPromise({
            try: () =>
              api.submitTurnCommand({
                sessionId: runtime,
                modelId,
                thinkingLevel,
                toolAccess,
                message,
                cwd: cwd.trim() || undefined,
                piSessionId,
                mode,
                queueAction,
                queueReplacement: contextualQueueReplacement,
                browserToolEnabled: browserEnabledForTurn,
                browserSessionId: runtime,
                browserBackend,
                skills,
                promptTemplates,
              }),
            catch: (error) => error,
          });
          updateSession(sessionId, (session) => ({
            ...session,
            piSessionId: result.piSessionId || session.piSessionId,
            contextUsage: api.runtimeContextUsage(result.status, session.contextUsage),
            status: "running",
          }));
          // Same acceptance bookkeeping the prompt path does. Without it a
          // steer/follow-up got no accept-grace (so a stale runtime-list
          // snapshot could idle the session and tear down its stream
          // mid-turn) and no cursor rewind if the runtime's seq had restarted.
          sessionRuntimeController().noteTurnAccepted(
            sessionId,
            undefined,
            result.status?.eventSeq,
          );
          if (result.piSessionId) onPiSessionIdChange?.(result.piSessionId);
          return { ok: true };
        }).pipe(
          Effect.catch((error) =>
            Effect.succeed({
              ok: false,
              error: error instanceof Error ? error.message : "Message failed",
            }),
          ),
        ),
      );
    },
    [
      browserToolEnabled,
      browserBackend,
      cwd,
      modelId,
      thinkingLevel,
      toolAccess,
      onPiSessionIdChange,
      updateSession,
    ],
  );

  const submitPrompt = useCallback(
    (args: SubmitArgs) =>
      submitPromptTurn(
        {
          activeTabId,
          browserToolEnabled,
          browserBackend,
          cwd,
          modelId,
          thinkingLevel,
          toolAccess,
          onPiSessionIdChange,
          selectionFor: selectionForRef.current,
          tabsRef,
          updateSession,
        },
        args,
      ),
    [
      activeTabId,
      modelId,
      thinkingLevel,
      toolAccess,
      cwd,
      browserToolEnabled,
      browserBackend,
      onPiSessionIdChange,
      updateSession,
    ],
  );

  const abortTurn = useCallback(
    (sessionId: SessionId) =>
      Effect.runPromise(
        Effect.gen(function* () {
          // Abort by the CONNECTION key (override-aware): after a server
          // restart the session's runtime lives under a different server key,
          // and /abort has no piSessionId fallback lookup.
          const runtime = sessionRuntimeController().connectionKey(sessionId);
          updateSession(sessionId, (session) => ({ ...session, status: "stopping" }));
          const cleared = yield* Effect.tryPromise({
            try: () => api.abortSession(runtime),
            catch: (error) => error,
          });
          // Settle the session fully. A direct status write bypasses the reducer
          // that normally finalizes tool badges on agent_end, and idling the
          // session detaches the SSE — so if the runtime's terminal event never
          // lands, any in-flight tool would render a perpetual "running" badge
          // and activeAssistantId would linger. Flush pending deltas first so the
          // last streamed text is committed before we finalize.
          sessionRuntimeController().flush(sessionId);
          updateSession(sessionId, settleTurnFinalizingTools);
          return cleared;
        }),
      ),
    [updateSession],
  );

  const loadAndReplay = useCallback(
    (piSessionId: string, sessionId: SessionId) => {
      if (inFlightReplays.has(sessionId)) return Promise.resolve();
      inFlightReplays.add(sessionId);
      return Effect.runPromise(
        Effect.gen(function* () {
          const cachedMessages = readTranscriptSnapshot(piSessionId);
          const seedCached = (session: Session) =>
            session.messages.length === 0 && cachedMessages
              ? { ...session, messages: cachedMessages }
              : session;
          if (!cwd) {
            updateSession(sessionId, (session) =>
              seedCached(session.status === "loading" ? { ...session, status: "idle" } : session),
            );
            return;
          }
          updateSession(sessionId, (session) => ({
            ...seedCached(session),
            status: "loading",
            error: "",
          }));
          // Canonical replay and the runtime-status probe are independent — the
          // status key is derived synchronously — so run them concurrently
          // instead of blocking the (now tail-limited) canonical read on the
          // status round-trip.
          const runtimeId = sessionRuntimeController().connectionKey(sessionId);
          const [replayResult, runtimeStatus] = yield* Effect.all(
            [
              Effect.tryPromise({
                try: () => api.loadCanonicalSession(piSessionId, cwd),
                catch: (error) => error,
              }).pipe(Effect.result),
              Effect.tryPromise({
                try: () => api.loadRuntimeStatus(runtimeId, piSessionId),
                catch: () => null,
              }),
            ],
            { concurrency: "unbounded" },
          );
          if (replayResult._tag === "Success") {
            const { events, cursor, meta } = replayResult.success;
            const runtimeActive = runtimeCanHydrateCanonicalSession(runtimeStatus, piSessionId);
            const replayEvents = mergeCanonicalAndRuntimeEvents(
              events,
              runtimeActive ? runtimeStatus?.events : [],
            );
            const {
              messages,
              title,
              startedAt,
              modelId: replayModelId,
              tokenStats,
            } = foldSessionEvents(replayEvents);
            const replaySeq = replayCursorAfterRuntimeHydration(runtimeStatus, piSessionId);
            updateSession(sessionId, (session) => ({
              ...session,
              messages: reconcileReplayMessages(session.messages, messages),
              piSessionId,
              cwd: session.cwd || cwd,
              // Head-scan meta carries the real session model/title; the fold's
              // own title would be the tail slice's first user message, not the
              // session's first prompt.
              modelId:
                session.modelId ||
                meta?.modelId ||
                replayModelId ||
                runtimeStatus?.modelId ||
                modelId,
              title: meta?.title ?? title ?? session.title,
              startedAt: meta?.startedAt ?? startedAt ?? session.startedAt,
              tokenStats: tokenStats ?? undefined,
              // Lifetime spend is computed server-side from the whole rollout,
              // so it survives both compaction and the tail load's cutoff.
              usageTotals: meta?.usage ?? session.usageTotals,
              contextUsage: api.runtimeContextUsage(runtimeStatus, session.contextUsage),
              status: runtimeActive ? "running" : "idle",
              activeAssistantId: undefined,
              // A non-null cursor means the tail load left older history unread;
              // the timeline shows a "Load earlier" affordance while it is set.
              historyCursor: messages.length > 0 ? cursor : (session.historyCursor ?? null),
              // The replay has landed, so whatever came from the snapshot has
              // been superseded and must not keep asking to be replayed.
              hydratedFromCache: false,
              error: "",
            }));
            // Reattach the live stream from the hydrated cursor so EventSource
            // does not replay already-rendered content.
            sessionRuntimeController().noteReplayHydrated(sessionId, replaySeq);
          } else {
            const err = replayResult.failure;
            // Canonical read failed. If the runtime is still alive, don't strand the
            // session idle (which would drop the live stream — reconcile only
            // subscribes for live statuses): keep the seeded history, mark it running,
            // and reset the cursor so the reattached SSE replays the runtime backlog.
            if (runtimeCanHydrateCanonicalSession(runtimeStatus, piSessionId)) {
              updateSession(sessionId, (session) => ({
                ...session,
                contextUsage: api.runtimeContextUsage(runtimeStatus, session.contextUsage),
                status: "running",
                activeAssistantId: undefined,
                error: "",
              }));
              sessionRuntimeController().noteReplayHydrated(sessionId, undefined);
              return;
            }
            updateSession(sessionId, (session) => ({
              ...session,
              error: err instanceof Error ? err.message : "Failed to load session",
              status: "idle",
            }));
          }
        }).pipe(
          Effect.ensuring(
            Effect.sync(() => {
              inFlightReplays.delete(sessionId);
            }),
          ),
        ),
      );
    },
    [cwd, modelId, updateSession],
  );

  // Page the previous (older) chunk of a tail-loaded transcript into view and
  // prepend it. Each page is snapped to a user-turn boundary and abuts the
  // current first message exactly (cursor = first loaded byte), so folding the
  // page on its own and prepending is equivalent to a single larger fold.
  const loadEarlier = useCallback(
    (sessionId: SessionId): Promise<void> => {
      const session = tabsRef.current.find((tab) => tab.id === sessionId);
      const cursor = session?.historyCursor;
      if (!session || !session.piSessionId || !cwd || cursor == null) return Promise.resolve();
      if (loadingEarlierRef.current.has(sessionId)) return Promise.resolve();
      loadingEarlierRef.current.add(sessionId);
      const piSessionId = session.piSessionId;
      return Effect.runPromise(
        Effect.gen(function* () {
          const result = yield* Effect.tryPromise({
            try: () => api.loadCanonicalSession(piSessionId, cwd, { before: cursor }),
            catch: (error) => error,
          }).pipe(Effect.result);
          if (result._tag !== "Success") return;
          const { messages: earlier } = foldSessionEvents(result.success.events);
          updateSession(sessionId, (current) => ({
            ...current,
            messages: earlier.length > 0 ? [...earlier, ...current.messages] : current.messages,
            historyCursor: result.success.cursor,
          }));
        }).pipe(
          Effect.ensuring(
            Effect.sync(() => {
              loadingEarlierRef.current.delete(sessionId);
            }),
          ),
        ),
      );
    },
    [cwd, updateSession],
  );

  const compact = useCallback(
    (sessionId: SessionId) =>
      Effect.runPromise(
        Effect.gen(function* () {
          const session = tabsRef.current.find((tab) => tab.id === sessionId);
          if (!session || !modelId) return;
          updateSession(sessionId, (s) => ({ ...s, error: "" }));
          const result = yield* Effect.tryPromise({
            try: () =>
              api.compactSession({
                sessionId: session.id,
                modelId,
                thinkingLevel,
                toolAccess,
                cwd: cwd.trim() || undefined,
                piSessionId: session.piSessionId,
                browserToolEnabled,
                browserSessionId: session.id,
                browserBackend,
                skills: selectionForRef.current(sessionId).skills ?? EMPTY_SKILLS,
                promptTemplates:
                  selectionForRef.current(sessionId).promptTemplates ?? EMPTY_PROMPT_TEMPLATES,
              }),
            catch: (error) => error,
          });
          const nextSessionId = result.status?.piSessionId || session.piSessionId;
          if (nextSessionId) {
            yield* Effect.tryPromise({
              try: () => loadAndReplay(nextSessionId, sessionId),
              catch: (error) => error,
            });
          }
          updateSession(sessionId, (s) => ({
            ...s,
            contextUsage: api.runtimeContextUsage(result.status ?? null, null),
            tokenStats: undefined,
          }));
        }).pipe(
          Effect.catch((error) =>
            Effect.sync(() => {
              updateSession(sessionId, (s) => ({
                ...s,
                error: error instanceof Error ? error.message : "Compaction failed",
              }));
            }),
          ),
        ),
      ),
    [browserToolEnabled, browserBackend, cwd, loadAndReplay, modelId, thinkingLevel, updateSession],
  );

  const acceptsControl = useCallback(
    async (
      tab: { status: Session["status"]; piSessionId?: string | null },
      runtime: string,
    ): Promise<boolean> => {
      // "stopping" counts: the composer still draws itself as running there, and
      // a turn being torn down can still take a follow-up for the next one.
      if (tab.status !== "running" && tab.status !== "starting" && tab.status !== "stopping") {
        return false;
      }
      const status = await loadRuntimeStatusCb(runtime, tab.piSessionId).catch(() => null);
      return runtimeStatusAcceptsControl(status, tab.piSessionId);
    },
    [loadRuntimeStatusCb],
  );

  return useMemo<SessionEngine>(
    () => ({
      submitPrompt,
      sendControl,
      loadRuntimeStatus: loadRuntimeStatusCb,
      abortTurn,
      loadAndReplay,
      loadEarlier,
      compact,
      acceptsControl,
    }),
    [
      submitPrompt,
      sendControl,
      loadRuntimeStatusCb,
      abortTurn,
      loadAndReplay,
      loadEarlier,
      compact,
      acceptsControl,
    ],
  );
}