import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import { truncateToWidth, visibleWidth, type Component } from "@earendil-works/pi-tui";
import { projectJobStatus, type JobStatus } from "./job-status.js";
import type { SubagentSessionSnapshot } from "./subagent-manager.js";
import type { WorkState } from "./types.js";

export interface LiveWidgetTheme {
  fg(color: "accent" | "dim" | "error" | "muted" | "success", text: string): string;
  bold(text: string): string;
}

export interface LiveWidgetRenderOptions {
  readonly now: number;
  readonly frame: number;
  readonly width: number;
  readonly theme: LiveWidgetTheme;
}

const SPINNER = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"] as const;
const LINGER_MS = 5_000;

const isTerminal = (state: WorkState): boolean => state === "completed" || state === "failed" || state === "cancelled";

type SessionRowKind = "opening" | "running" | "cancelling" | "closing" | "queued" | "follow_up" | "waiting" | "completed" | "failed" | "cancelled";
type TerminalRowKind = Extract<SessionRowKind, "completed" | "failed" | "cancelled">;

const isTerminalRowKind = (kind: SessionRowKind): kind is TerminalRowKind =>
  kind === "completed" || kind === "failed" || kind === "cancelled";

const isLingering = (session: Readonly<SubagentSessionSnapshot>, now: number): boolean =>
  isTerminal(session.generation.state)
  && session.generation.finishedAt !== undefined
  && now < session.generation.finishedAt + LINGER_MS;

interface SessionRow {
  session: SubagentSessionSnapshot;
  status: JobStatus;
  kind: SessionRowKind;
}

interface DisplaySessions {
  rows: readonly SessionRow[];
  idleOpenCount: number;
}

const rowKind = (session: Readonly<SubagentSessionSnapshot>, now: number): SessionRowKind | "idle" | undefined => {
  if (session.state === "closed") return undefined;
  if (session.state === "opening") return "opening";
  if (session.state === "closing") return "closing";

  const workState = session.generation.state;
  if (workState === "running") return "running";
  if (workState === "cancelling") return "cancelling";
  if (workState === "queued") return "queued";
  if (workState === "waiting_for_parent") return "waiting";
  if (session.state === "open" && session.queuedFollowUp) return "follow_up";
  if (isLingering(session, now)) return workState;
  return session.state === "open" ? "idle" : undefined;
};

const displaySessions = (sessions: readonly SubagentSessionSnapshot[], now: number): DisplaySessions => {
  const groups: Record<Exclude<SessionRowKind, "queued" | "follow_up" | "completed" | "failed" | "cancelled">, SessionRow[]> = {
    opening: [], running: [], cancelling: [], closing: [], waiting: [],
  };
  const queuedRows: SessionRow[] = [];
  const terminalRows: SessionRow[] = [];
  let idleOpenCount = 0;

  for (const session of sessions) {
    const kind = rowKind(session, now);
    if (kind === "idle") {
      idleOpenCount += 1;
      continue;
    }
    if (kind === undefined) continue;
    const row = { session, status: projectJobStatus(session, now), kind };
    if (kind === "queued" || kind === "follow_up") queuedRows.push(row);
    else if (isTerminalRowKind(kind)) terminalRows.push(row);
    else groups[kind].push(row);
  }

  return {
    rows: [
      ...groups.opening,
      ...groups.running,
      ...groups.cancelling,
      ...groups.closing,
      ...queuedRows,
      ...groups.waiting,
      ...terminalRows,
    ],
    idleOpenCount,
  };
};

const formatDuration = (milliseconds: number): string => `${(Math.max(0, milliseconds) / 1_000).toFixed(1)}s`;

const rowLabel = (kind: SessionRowKind): string => {
  if (kind === "follow_up") return "follow-up queued";
  if (kind === "waiting") return "waiting for parent";
  return kind;
};

const stateIcon = (kind: SessionRowKind, frame: number, theme: LiveWidgetTheme): string => {
  if (kind === "running") {
    const spinner = SPINNER[((frame % SPINNER.length) + SPINNER.length) % SPINNER.length];
    if (spinner === undefined) return theme.fg("accent", SPINNER[0]);
    return theme.fg("accent", spinner);
  }
  if (kind === "opening" || kind === "cancelling" || kind === "closing") return theme.fg("accent", "◌");
  if (kind === "queued" || kind === "follow_up") return theme.fg("muted", "○");
  if (kind === "waiting") return theme.fg("error", "!");
  if (kind === "completed") return theme.fg("success", "✓");
  if (kind === "failed") return theme.fg("error", "✗");
  return theme.fg("dim", "■");
};

const formatTokens = (count: number): string => {
  const safe = Math.max(0, count);
  if (safe >= 1_000_000) return `${(safe / 1_000_000).toFixed(1)}M tokens`;
  if (safe >= 1_000) return `${(safe / 1_000).toFixed(1)}k tokens`;
  return `${safe} token${safe === 1 ? "" : "s"}`;
};

const compactThinking = (thinking: string | undefined): string | undefined => thinking?.replace(/\s+\([^)]*\)$/u, "");

const formatStats = (status: Readonly<JobStatus>): string => {
  const parts: string[] = [];
  if (status.usage.turns > 0) parts.push(`↻${status.usage.turns}`);
  const tokens = Math.max(0, status.usage.input + status.usage.output);
  if (tokens > 0) parts.push(formatTokens(tokens));
  const selectedModel = [status.reportedModel, status.launchModel].find((model) => Boolean(model));
  const model = selectedModel?.split("/").at(-1);
  if (model) parts.push(model);
  const thinking = compactThinking(status.launchThinking);
  if (thinking) parts.push(thinking);
  const duration = status.workState === "queued" ? status.queueDurationMs : status.runDurationMs ?? status.queueDurationMs;
  if (duration !== undefined) parts.push(status.workState === "queued" ? `queued ${formatDuration(duration)}` : formatDuration(duration));
  return parts.join(" · ");
};

const formatSessionRow = (prefix: string, facts: string, details: string, activity: string, width: number): string => {
  const safeWidth = Math.max(0, width);
  const prefixWidth = visibleWidth(prefix);
  if (prefixWidth >= safeWidth) return truncateToWidth(prefix, safeWidth, "");

  const core = `${prefix}  ${facts}`;
  if (visibleWidth(core) >= safeWidth) return truncateToWidth(core, safeWidth);

  const coreWithDetails = details ? `${core} ${details}` : core;
  if (visibleWidth(coreWithDetails) > safeWidth) return core;
  if (!activity) return coreWithDetails;

  const activityWidth = safeWidth - visibleWidth(coreWithDetails) - 1;
  return activityWidth >= 5 ? `${coreWithDetails} ${truncateToWidth(activity, activityWidth)}` : coreWithDetails;
};

const idleSummary = (count: number): string =>
  `${count} idle subagent session${count === 1 ? " remains" : "s remain"} open`;

export function formatLiveWidgetLines(sessions: readonly SubagentSessionSnapshot[], options: LiveWidgetRenderOptions): string[] {
  const { now, frame, width, theme } = options;
  const visible = displaySessions(sessions, now);
  if (visible.rows.length === 0 && visible.idleOpenCount === 0) return [];

  const hasActionable = visible.rows.some((row) => !isTerminalRowKind(row.kind));
  const lines = [theme.fg(hasActionable ? "accent" : "dim", `${hasActionable ? "●" : "○"} Subagents`)];

  visible.rows.forEach((row, index) => {
    const isLast = index === visible.rows.length - 1;
    const connector = isLast ? "└─" : "├─";
    const agent = row.kind === "running" ? theme.bold(row.status.agent) : theme.fg("dim", row.status.agent);
    const prefix = `${theme.fg("dim", connector)} ${stateIcon(row.kind, frame, theme)} ${agent}`;
    const facts = theme.fg("muted", `${row.status.id} · gen ${row.status.generationNumber ?? 0} · ${rowLabel(row.kind)}`);
    const stats = formatStats(row.status);
    const details = stats ? theme.fg("dim", `· ${stats}`) : "";
    const summary = isTerminalRowKind(row.kind) ? undefined : row.status.recentActivity.at(-1)?.summary;
    const activity = summary ? theme.fg("dim", `· ${summary}`) : "";
    lines.push(formatSessionRow(prefix, facts, details, activity, width));
  });

  if (visible.idleOpenCount > 0) lines.push(theme.fg("dim", idleSummary(visible.idleOpenCount)));
  return lines.map((line) => truncateToWidth(line, Math.max(0, width)));
}

type TimerHandle = unknown;
type WidgetUi = Pick<ExtensionContext["ui"], "setWidget">;

export interface LiveSubagentsWidgetOptions {
  readonly now?: () => number;
  readonly setInterval?: (callback: () => void, delay: number) => TimerHandle;
  readonly clearInterval?: (handle: TimerHandle) => void;
  readonly setTimeout?: (callback: () => void, delay: number) => TimerHandle;
  readonly clearTimeout?: (handle: TimerHandle) => void;
}

export class LiveSubagentsWidget {
  private readonly now: () => number;
  private readonly startInterval: (callback: () => void, delay: number) => TimerHandle;
  private readonly stopInterval: (handle: TimerHandle) => void;
  private readonly startTimeout: (callback: () => void, delay: number) => TimerHandle;
  private readonly stopTimeout: (handle: TimerHandle) => void;
  private ui: WidgetUi | undefined;
  private sessions: readonly SubagentSessionSnapshot[] = [];
  private requestRender: (() => void) | undefined;
  private animation: TimerHandle | undefined;
  private expiry: TimerHandle | undefined;
  private expiryDeadline: number | undefined;
  private frame = 0;
  private registered = false;
  private disposed = false;

  constructor(options: LiveSubagentsWidgetOptions = {}) {
    this.now = options.now ?? Date.now;
    this.startInterval = options.setInterval ?? ((callback, delay) => setInterval(callback, delay));
    this.stopInterval = options.clearInterval ?? ((handle) => clearInterval(handle as ReturnType<typeof setInterval>));
    this.startTimeout = options.setTimeout ?? ((callback, delay) => setTimeout(callback, delay));
    this.stopTimeout = options.clearTimeout ?? ((handle) => clearTimeout(handle as ReturnType<typeof setTimeout>));
  }

  setUi(ui: WidgetUi): void {
    if (this.disposed || this.ui === ui) return;
    this.clearRegistration();
    this.ui = ui;
    this.refresh(false);
  }

  setSessions(sessions: readonly SubagentSessionSnapshot[]): void {
    if (this.disposed) return;
    this.sessions = sessions;
    this.refresh(true);
  }

  dispose(): void {
    if (this.disposed) return;
    this.disposed = true;
    this.stopAnimation();
    this.stopExpiry();
    this.clearRegistration();
    this.ui = undefined;
    this.sessions = [];
  }

  private refresh(fromSessionUpdate: boolean): void {
    if (this.disposed || !this.ui) return;
    const now = this.now();
    const visible = displaySessions(this.sessions, now);
    if (visible.rows.length === 0 && visible.idleOpenCount === 0) {
      this.stopAnimation();
      this.stopExpiry();
      this.clearRegistration();
      return;
    }

    const wasAnimating = this.animation !== undefined;
    this.ensureRegistration();
    if (visible.rows.some((row) => row.kind === "running")) this.ensureAnimation();
    else this.stopAnimation();
    this.syncExpiry(visible.rows, now);
    if (fromSessionUpdate && (!wasAnimating || this.animation === undefined)) this.requestRender?.();
  }

  private ensureRegistration(): void {
    if (this.registered || !this.ui) return;
    this.ui.setWidget("simple-subagents", (tui, theme): Component => {
      this.requestRender = () => tui.requestRender();
      return {
        render: (width) => formatLiveWidgetLines(this.sessions, {
          now: this.now(), frame: this.frame, width, theme: theme as LiveWidgetTheme,
        }),
        invalidate: () => {},
      };
    }, { placement: "aboveEditor" });
    this.registered = true;
  }

  private ensureAnimation(): void {
    if (this.animation !== undefined) return;
    this.animation = this.startInterval(() => {
      if (this.disposed) return;
      this.frame = (this.frame + 1) % SPINNER.length;
      this.requestRender?.();
    }, 80);
  }

  private stopAnimation(): void {
    if (this.animation === undefined) return;
    this.stopInterval(this.animation);
    this.animation = undefined;
  }

  private syncExpiry(rows: readonly SessionRow[], now: number): void {
    const deadlines = rows
      .filter((row) => isTerminalRowKind(row.kind))
      .flatMap((row) => row.session.generation.finishedAt === undefined ? [] : [row.session.generation.finishedAt + LINGER_MS]);
    const deadline = deadlines.length === 0 ? undefined : Math.min(...deadlines);
    if (deadline === this.expiryDeadline) return;
    this.stopExpiry();
    if (deadline === undefined) return;
    this.expiryDeadline = deadline;
    this.expiry = this.startTimeout(() => {
      if (this.disposed) return;
      this.expiry = undefined;
      this.expiryDeadline = undefined;
      this.refresh(false);
      this.requestRender?.();
    }, Math.max(0, deadline - now));
  }

  private stopExpiry(): void {
    if (this.expiry !== undefined) this.stopTimeout(this.expiry);
    this.expiry = undefined;
    this.expiryDeadline = undefined;
  }

  private clearRegistration(): void {
    this.requestRender = undefined;
    if (!this.registered) return;
    this.ui?.setWidget("simple-subagents", undefined);
    this.registered = false;
  }

}