import { type ChildProcess, spawn } from "node:child_process";
import { mkdtemp, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { dirname, join } from "node:path";
import { PI_TEXTUAL_LIB } from "./driver.ts";
import {
  buildSessionDriver,
  type SessionDriverConfig,
} from "./session-driver.ts";

/** Result envelope emitted by the session driver. */
export interface SessionResult {
  appExitValue?: unknown;
  closed?: boolean;
  error?: string;
  id?: number;
  kind?: string;
  size?: [number, number];
  svgBytes?: number;
  svgPath?: string;
  text?: string;
  tree?: unknown;
  widgets?: unknown;
  [key: string]: unknown;
}

/** A screen-text update streamed while an actions command runs. */
export interface SessionUpdate {
  index: number;
  text: string;
}

interface PendingCommand {
  id: number;
  onUpdate?: (update: SessionUpdate) => void;
  reject: (error: Error) => void;
  resolve: (result: SessionResult) => void;
}

export interface TextualSession {
  /** Absolute path of the app file that runs in this session. */
  appPath: string;
  /** The Python subprocess. */
  child: ChildProcess;
  /** True once the session can no longer serve commands. */
  closed: boolean;
  /** Timer that closes idle sessions (unref'd so it never holds pi open). */
  idleTimer: NodeJS.Timeout | null;
  /** Session name used by the agent to address the session. */
  name: string;
  nextId: number;
  /** The command currently waiting for a result line. */
  pending: PendingCommand | null;
  /** Python interpreter that runs the app. */
  python: string;
  /** Serializes commands on this session; one in flight at a time. */
  queue: Promise<unknown>;
  /** Partial stdout line accumulation. */
  stdoutBuffer: string;
}

export interface OpenSessionOptions {
  appPath: string;
  cwd: string;
  name: string;
  python: string;
  signal?: AbortSignal;
  size: [number, number];
  timeoutMs: number;
}

export interface CommandOptions {
  onUpdate?: (update: SessionUpdate) => void;
  signal?: AbortSignal;
  timeoutMs: number;
}

const IDLE_TIMEOUT_MS = 15 * 60_000;
const KILL_GRACE_MS = 1000;

/**
 * Manages persistent Textual sessions: one Python subprocess per session, with
 * the app kept alive and driven through commands on stdin.
 */
export class SessionManager {
  readonly #sessions = new Map<string, TextualSession>();
  readonly #dirs = new Set<string>();
  readonly #maxSessions: number;
  readonly #spawnImpl: typeof spawn;

  constructor(spawnImpl: typeof spawn = spawn, maxSessions = 8) {
    this.#spawnImpl = spawnImpl;
    this.#maxSessions = maxSessions;
  }

  /** Names of the currently open sessions. */
  names(): string[] {
    return [...this.#sessions.keys()];
  }

  isOpen(name: string): boolean {
    return this.#sessions.has(name);
  }

  /** App path of an open session, for path bookkeeping on screenshots. */
  appPathOf(name: string): string | undefined {
    return this.#sessions.get(name)?.appPath;
  }

  /** Opens a session: spawns the driver and waits until the app is ready. */
  async open(options: OpenSessionOptions): Promise<TextualSession> {
    if (this.#sessions.has(options.name)) {
      throw new Error(
        `textual_run: session '${options.name}' is already open. Drop 'app' to act on it.`
      );
    }
    if (this.#sessions.size >= this.#maxSessions) {
      throw new Error(
        `textual_run: too many open sessions (${this.#maxSessions}); close one first: ${this.names().join(", ")}`
      );
    }
    const dir = await mkdtemp(join(tmpdir(), "pi-textual-session-"));
    this.#dirs.add(dir);
    const config: SessionDriverConfig = {
      appPath: options.appPath,
      size: options.size,
    };
    const driverPath = join(dir, "session-driver.py");
    await writeFile(driverPath, buildSessionDriver(config));
    await writeFile(join(dir, "pi_textual_lib.py"), PI_TEXTUAL_LIB);
    const child = this.#spawnImpl(options.python, [driverPath], {
      cwd: options.cwd,
      stdio: ["pipe", "pipe", "pipe"],
    });
    const session: TextualSession = {
      appPath: options.appPath,
      child,
      closed: false,
      idleTimer: null,
      name: options.name,
      nextId: 0,
      pending: null,
      python: options.python,
      queue: Promise.resolve(),
      stdoutBuffer: "",
    };
    this.#sessions.set(options.name, session);
    this.#attachChild(session);
    this.#refreshIdle(session);
    try {
      const result = await this.#waitFor(
        session,
        options.timeoutMs,
        options.signal
      );
      if (result.closed) {
        this.#kill(session);
        const exit =
          "appExitValue" in result
            ? ` (exit value: ${JSON.stringify(result.appExitValue)})`
            : "";
        throw new Error(
          `textual_run: the app exited before session '${options.name}' was ready${exit}`
        );
      }
      return session;
    } catch (error) {
      this.#kill(session);
      throw error;
    }
  }

  /**
   * Sends one command to the session and resolves with its result line.
   * Commands on the same session run strictly in order.
   */
  command(
    name: string,
    command: Record<string, unknown>,
    options: CommandOptions
  ): Promise<SessionResult> {
    const session = this.#sessions.get(name);
    if (!session || session.closed) {
      throw new Error(
        `textual_run: no open session '${name}' (${this.#openOrClosed()}). Pass 'app' with 'session' to open it, or omit 'session' for a one-shot run.`
      );
    }
    this.#refreshIdle(session);
    const run = () => this.#runCommand(session, command, options);
    const queued = session.queue.then(run, run);
    session.queue = queued;
    return queued;
  }

  /**
   * Closes a session: asks the driver to exit, then reaps the process. A
   * session the driver already declared closed is still reaped, so a wedged
   * child process cannot outlive the close call.
   */
  async close(name: string, options: CommandOptions): Promise<SessionResult> {
    const session = this.#sessions.get(name);
    if (!session) {
      return { closed: true };
    }
    this.#refreshIdle(session);
    let result: SessionResult;
    try {
      result = await this.command(name, { command: "close" }, options);
    } catch {
      result = { closed: true };
    }
    this.#kill(session);
    return result;
  }

  /** Closes every open session; called on pi session shutdown. */
  shutdownAll(): void {
    for (const session of this.#sessions.values()) {
      this.#kill(session);
    }
    this.#sessions.clear();
  }

  #openOrClosed(): string {
    const open = this.names();
    return open.length > 0 ? `open: ${open.join(", ")}` : "no sessions open";
  }

  #refreshIdle(session: TextualSession): void {
    if (session.idleTimer) {
      clearTimeout(session.idleTimer);
    }
    session.idleTimer = setTimeout(() => {
      if (!session.closed) {
        this.#kill(session);
      }
    }, IDLE_TIMEOUT_MS);
    session.idleTimer.unref();
  }

  #attachChild(session: TextualSession): void {
    let stderrTail = "";
    session.child.stderr?.on("data", (chunk: Buffer) => {
      stderrTail = (stderrTail + chunk.toString()).slice(-2000);
    });
    session.child.stdout?.on("data", (chunk: Buffer) => {
      session.stdoutBuffer += chunk.toString();
      let newline = session.stdoutBuffer.indexOf("\n");
      while (newline >= 0) {
        const line = session.stdoutBuffer.slice(0, newline).trim();
        session.stdoutBuffer = session.stdoutBuffer.slice(newline + 1);
        if (line) {
          this.#handleLine(session, line);
        }
        newline = session.stdoutBuffer.indexOf("\n");
      }
    });
    session.child.on("error", (error) => {
      this.#failPending(
        session,
        new Error(
          `textual_run: session '${session.name}' could not start: ${error.message}`
        )
      );
    });
    session.child.on("exit", (code, signal) => {
      session.closed = true;
      this.#sessions.delete(session.name);
      const stderr = stderrTail ? `\nstderr: ${stderrTail}` : "";
      this.#failPending(
        session,
        new Error(
          `textual_run: session '${session.name}' ended (exit ${code ?? "unknown"}${signal ? `, ${signal}` : ""})${stderr}`
        )
      );
      const dir = this.#dirOf(session);
      if (dir) {
        this.#dirs.delete(dir);
        rm(dir, { force: true, recursive: true }).catch(() => {
          // Best effort: the directory lives in the system temp dir.
        });
      }
    });
  }

  #dirOf(session: TextualSession): string | undefined {
    // The driver file lives in the session's temp dir; find it via the
    // child's spawn args. Kept simple: the driver path is args[1].
    const args = session.child.spawnargs ?? [];
    const driverIndex = args.findIndex((arg) =>
      arg.endsWith("session-driver.py")
    );
    return driverIndex >= 0 ? dirname(args[driverIndex]) : undefined;
  }

  #handleLine(session: TextualSession, line: string): void {
    let doc: SessionResult;
    try {
      doc = JSON.parse(line) as SessionResult;
    } catch {
      // Not a driver document; ignore (the driver prints nothing else, but
      // a misbehaving app could). Keep the session alive.
      return;
    }
    if (doc.kind === "update") {
      session.pending?.onUpdate?.({
        index: typeof doc.index === "number" ? doc.index : 0,
        text: doc.text ?? "",
      });
      return;
    }
    const { pending } = session;
    if (!pending) {
      // Unsolicited result: the driver's final message after the app exited
      // on its own. Mark the session closed and reap the child now; if the
      // driver fails to exit (a wedged app), killing it is the only way the
      // parent does not wait on the process forever.
      session.closed = true;
      this.#kill(session);
      return;
    }
    if (doc.id !== undefined && doc.id !== pending.id) {
      // Result for a command we no longer track; keep waiting for ours.
      return;
    }
    session.pending = null;
    if (doc.error) {
      pending.reject(
        new Error(`textual_run: the app failed: ${doc.error.slice(0, 4000)}`)
      );
      return;
    }
    if (doc.closed) {
      // The app exited; the session can no longer serve commands.
      session.closed = true;
      pending.resolve(doc);
      return;
    }
    pending.resolve(doc);
  }

  #waitFor(
    session: TextualSession,
    timeoutMs: number,
    signal?: AbortSignal
  ): Promise<SessionResult> {
    return new Promise((resolve, reject) => {
      const timer = setTimeout(() => {
        this.#failPending(
          session,
          new Error(
            `textual_run: session '${session.name}' did not become ready within ${timeoutMs} ms`
          )
        );
        reject(
          new Error(
            `textual_run: session '${session.name}' did not become ready within ${timeoutMs} ms`
          )
        );
      }, timeoutMs);
      timer.unref();
      const onAbort = () => {
        clearTimeout(timer);
        this.#failPending(
          session,
          new Error(`textual_run: session '${session.name}' open was aborted`)
        );
        reject(
          new Error(`textual_run: session '${session.name}' open was aborted`)
        );
      };
      signal?.addEventListener("abort", onAbort, { once: true });
      session.pending = {
        id: -1,
        reject: (error) => {
          clearTimeout(timer);
          signal?.removeEventListener("abort", onAbort);
          reject(error);
        },
        resolve: (result) => {
          clearTimeout(timer);
          signal?.removeEventListener("abort", onAbort);
          resolve(result);
        },
      };
    });
  }

  #runCommand(
    session: TextualSession,
    command: Record<string, unknown>,
    options: CommandOptions
  ): Promise<SessionResult> {
    return new Promise((resolve, reject) => {
      if (session.closed) {
        reject(new Error(`textual_run: session '${session.name}' is closed`));
        return;
      }
      session.nextId += 1;
      const id = session.nextId;
      let settled = false;
      const timer = setTimeout(() => {
        if (settled) {
          return;
        }
        settled = true;
        this.#failPending(
          session,
          new Error(
            `textual_run: session '${session.name}' timed out after ${options.timeoutMs} ms and was closed`
          )
        );
        this.#kill(session);
        reject(new Error(`textual_run: session '${session.name}' timed out`));
      }, options.timeoutMs);
      timer.unref();
      const onAbort = () => {
        if (settled) {
          return;
        }
        settled = true;
        clearTimeout(timer);
        this.#failPending(
          session,
          new Error(
            `textual_run: session '${session.name}' was aborted and closed`
          )
        );
        this.#kill(session);
        reject(new Error(`textual_run: session '${session.name}' was aborted`));
      };
      if (options.signal?.aborted) {
        onAbort();
      } else {
        options.signal?.addEventListener("abort", onAbort, { once: true });
      }
      session.pending = {
        id,
        onUpdate: options.onUpdate,
        reject: (error) => {
          if (settled) {
            return;
          }
          settled = true;
          clearTimeout(timer);
          options.signal?.removeEventListener("abort", onAbort);
          reject(error);
        },
        resolve: (result) => {
          if (settled) {
            return;
          }
          settled = true;
          clearTimeout(timer);
          options.signal?.removeEventListener("abort", onAbort);
          resolve(result);
        },
      };
      session.child.stdin?.write(`${JSON.stringify({ ...command, id })}\n`);
    });
  }

  #failPending(session: TextualSession, error: Error): void {
    const { pending } = session;
    session.pending = null;
    pending?.reject(error);
  }

  #kill(session: TextualSession): void {
    if (session.idleTimer) {
      clearTimeout(session.idleTimer);
      session.idleTimer = null;
    }
    session.closed = true;
    this.#sessions.delete(session.name);
    const { child } = session;
    this.#failPending(
      session,
      new Error(`textual_run: session '${session.name}' was closed`)
    );
    // Idempotent: killing an already-dead child is a no-op, and a wedged
    // child gets SIGTERM then SIGKILL so the parent never waits on it.
    if (child.exitCode === null) {
      try {
        child.kill("SIGTERM");
      } catch {
        // Already gone.
      }
      const force = setTimeout(() => {
        try {
          child.kill("SIGKILL");
        } catch {
          // Already gone.
        }
      }, KILL_GRACE_MS);
      force.unref();
    }
  }
}