import { spawn, type ChildProcess } from "node:child_process";
import { createWriteStream } from "node:fs";
import { mkdtemp, rm, stat, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { StringDecoder } from "node:string_decoder";
import { isAbsolute, join, resolve as resolvePath } from "node:path";
import { keyText, type ExtensionAPI } from "@earendil-works/pi-coding-agent";
import {
  stripTerminalSequences,
  Text,
  truncateToWidth,
  type Component,
  visibleWidth,
  wrapTextWithAnsi,
} from "@earendil-works/pi-tui";
import { Type } from "typebox";

type Discovery = { exe: string; shell: "pwsh"; version: string };

export type ExecResult = {
  shell: string;
  shellVersion: string;
  shellExe: string;
  cwd: string;
  command: string;
  exitCode: number;
  stdout: string;
  stderr: string;
  timedOut: boolean;
  aborted: boolean;
  truncated: boolean;
  stdoutFullPath?: string;
  stderrFullPath?: string;
};

type CaptureResult = {
  text: string;
  truncated: boolean;
  fullPath?: string;
};

const MAX_ENCODED_COMMAND_ARG_LENGTH = 24_000;
const DEFAULT_STDOUT_MAX_BYTES = 20_000;
const DEFAULT_STDOUT_MAX_LINES = 500;
const STDERR_MAX_BYTES = 8_000;
const STDERR_MAX_LINES = 120;
const COLLAPSED_COMMAND_MAX_LENGTH = 120;
const COLLAPSED_OUTPUT_MAX_LINES = 3;

let cachedDiscovery: Discovery | null = null;

function clampInteger(value: number | undefined, fallback: number, max: number) {
  if (!Number.isFinite(value) || value === undefined) return fallback;
  return Math.min(Math.max(Math.floor(value), 1), max);
}

function formatBytes(bytes: number) {
  if (bytes < 1024) return `${bytes}B`;
  const kb = bytes / 1024;
  if (kb < 1024) return `${Math.round(kb)}KB`;
  return `${(kb / 1024).toFixed(1)}MB`;
}

function takeLastUtf8Bytes(text: string, maxBytes: number) {
  let out = "";
  let bytes = 0;

  for (let i = text.length - 1; i >= 0; i--) {
    const ch = text[i];
    const chBytes = Buffer.byteLength(ch, "utf-8");
    if (bytes + chBytes > maxBytes) break;
    out = ch + out;
    bytes += chBytes;
  }

  return out;
}

export function truncateTail(text: string, maxBytes: number, maxLines: number) {
  let truncated = false;
  let output = text;

  const lines = output.split(/\r?\n/);
  if (lines.length > maxLines) {
    output = lines.slice(-maxLines).join("\n");
    truncated = true;
  }

  if (Buffer.byteLength(output, "utf-8") > maxBytes) {
    output = takeLastUtf8Bytes(output, maxBytes);
    truncated = true;
  }

  return { text: output, truncated };
}

class OutputCapture {
  private text = "";
  private totalBytes = 0;
  private lineBreaks = 0;
  private sawContent = false;
  private truncatedInMemory = false;
  private closed = false;
  private writeError: Error | undefined;
  private readonly decoder = new StringDecoder("utf8");
  private readonly stream: ReturnType<typeof createWriteStream>;

  constructor(
    private readonly fullPath: string,
    private readonly maxBytes: number,
    private readonly maxLines: number
  ) {
    this.stream = createWriteStream(fullPath, { encoding: "utf8" });
    this.stream.on("error", (error) => {
      this.writeError = error;
    });
  }

  append(chunk: Buffer | string) {
    if (this.closed) return;

    const text = typeof chunk === "string" ? chunk : this.decoder.write(chunk);
    this.appendText(text);
  }

  private appendText(text: string) {
    if (!text) return;

    this.stream.write(text);
    this.sawContent = true;
    this.totalBytes += Buffer.byteLength(text, "utf-8");
    this.lineBreaks += text.match(/\n/g)?.length ?? 0;

    const truncated = truncateTail(this.text + text, this.maxBytes, this.maxLines);
    this.text = truncated.text;
    this.truncatedInMemory ||= truncated.truncated;
  }

  async finish(): Promise<CaptureResult> {
    this.appendText(this.decoder.end());
    this.closed = true;
    await new Promise<void>((resolve) => this.stream.end(resolve));

    const totalLines = this.sawContent ? this.lineBreaks + 1 : 0;
    const wasTruncated =
      this.truncatedInMemory || this.totalBytes > this.maxBytes || totalLines > this.maxLines;

    if (!wasTruncated) {
      return { text: this.text, truncated: false };
    }

    const outputLines = this.text ? this.text.split(/\r?\n/).length : 0;
    const outputBytes = Buffer.byteLength(this.text, "utf-8");
    const outputDescription =
      totalLines > outputLines
        ? `Showing lines ${totalLines - outputLines + 1}-${totalLines} of ${totalLines}`
        : `Showing last ${formatBytes(outputBytes)} of ${formatBytes(this.totalBytes)}`;
    const fullOutputNotice = this.writeError
      ? ` Full output could not be saved: ${this.writeError.message}]`
      : ` Full output: ${this.fullPath}]`;
    const notice = `\n\n[${outputDescription}.${fullOutputNotice}`;

    return {
      text: this.text.trimEnd() + notice,
      truncated: true,
      fullPath: this.writeError ? undefined : this.fullPath,
    };
  }
}

export function formatCommandPreview(command: string, expanded = true) {
  if (expanded) return command;

  const singleLine = command.replace(/\s+/g, " ").trim();
  return stripTerminalSequences(
    truncateToWidth(singleLine, COLLAPSED_COMMAND_MAX_LENGTH, "…"),
  );
}

function getTextContent(result: { content: Array<{ type: string; text?: string }> }) {
  return stripTerminalSequences(
    result.content
      .filter((content) => content.type === "text" && typeof content.text === "string")
      .map((content) => content.text ?? "")
      .join("\n"),
  );
}

class WidthAwareText implements Component {
  private readonly text = new Text("", 0, 0);
  private renderText: (width: number) => string = () => "";
  private cachedWidth: number | undefined;
  private cachedText: string | undefined;

  setRenderText(renderText: (width: number) => string) {
    this.renderText = renderText;
    this.cachedWidth = undefined;
    this.cachedText = undefined;
    this.text.invalidate();
  }

  render(width: number) {
    if (this.cachedWidth !== width) {
      const nextText = this.renderText(width);
      if (nextText !== this.cachedText) {
        this.text.setText(nextText);
        this.cachedText = nextText;
      }
      this.cachedWidth = width;
    }

    return this.text.render(width);
  }

  invalidate() {
    this.cachedWidth = undefined;
    this.cachedText = undefined;
    this.text.invalidate();
  }
}

export function formatResultSummary(details: Partial<ExecResult> | undefined) {
  if (!details) return "PowerShell finished";

  const markers = [
    typeof details.exitCode === "number" ? `exit ${details.exitCode}` : null,
    details.timedOut ? "timed out" : null,
    details.aborted ? "aborted" : null,
    details.truncated ? "truncated" : null,
  ].filter(Boolean);

  const shell = [details.shell, details.shellVersion === "unknown" ? undefined : details.shellVersion]
    .filter(Boolean)
    .join(" ");

  return [shell || "PowerShell", markers.length ? `(${markers.join(", ")})` : null]
    .filter(Boolean)
    .join(" ");
}

function trimOuterLineBreaks(text: string) {
  return text.replace(/^(?:\r?\n)+/, "").trimEnd();
}

export function formatOutputForContent(result: ExecResult, emptyText = "(no output)") {
  const stdout = trimOuterLineBreaks(result.stdout);
  const stderr = trimOuterLineBreaks(result.stderr);

  if (stdout && stderr) return `${stdout}\n\nSTDERR:\n${stderr}`;
  if (stdout) return stdout;
  if (stderr) return `STDERR:\n${stderr}`;
  return emptyText;
}

function appendStatus(text: string, status: string) {
  return text ? `${text}\n\n${status}` : status;
}

function commandWithPrelude(command: string) {
  const prelude = [
    "[Console]::OutputEncoding = [System.Text.UTF8Encoding]::new($false)",
    "$OutputEncoding = [Console]::OutputEncoding",
  ].join("; ");

  return `${prelude}\n${command}`;
}

function encodedCommand(command: string) {
  return Buffer.from(commandWithPrelude(command), "utf16le").toString("base64");
}

export function normalizeCwd(cwd: string | undefined, fallback: string) {
  if (!cwd) return fallback;
  const normalized = cwd.startsWith("@") ? cwd.slice(1) : cwd;
  return isAbsolute(normalized) ? normalized : resolvePath(fallback, normalized);
}

function encodedPowerShellArgs(command: string) {
  return [
    "-NoProfile",
    "-NonInteractive",
    "-ExecutionPolicy",
    "Bypass",
    "-OutputFormat",
    "Text",
    "-EncodedCommand",
    encodedCommand(command),
  ];
}

async function preparePowerShellArgs(command: string, tempDir: string) {
  const encoded = encodedCommand(command);
  if (encoded.length <= MAX_ENCODED_COMMAND_ARG_LENGTH) {
    return [
      "-NoProfile",
      "-NonInteractive",
      "-ExecutionPolicy",
      "Bypass",
      "-OutputFormat",
      "Text",
      "-EncodedCommand",
      encoded,
    ];
  }

  const scriptPath = join(tempDir, "command.ps1");
  await writeFile(scriptPath, commandWithPrelude(command), "utf8");
  return [
    "-NoProfile",
    "-NonInteractive",
    "-ExecutionPolicy",
    "Bypass",
    "-OutputFormat",
    "Text",
    "-File",
    scriptPath,
  ];
}

function killProcessTree(proc: ChildProcess) {
  if (proc.killed) return;

  if (process.platform === "win32" && proc.pid) {
    const killer = spawn("taskkill", ["/PID", String(proc.pid), "/T", "/F"], {
      windowsHide: true,
      stdio: "ignore",
    });
    killer.on("error", () => proc.kill());
    return;
  }

  proc.kill("SIGTERM");
}

export function parseMajorVersion(version: string) {
  const major = Number.parseInt(version.trim().split(/[.\s]/)[0] ?? "", 10);
  return Number.isFinite(major) ? major : undefined;
}

function probePwsh7(timeoutMs = 5_000) {
  return new Promise<Discovery | null>((resolve) => {
    const proc = spawn("pwsh", encodedPowerShellArgs("$PSVersionTable.PSVersion.ToString()"), {
      windowsHide: true,
      stdio: ["ignore", "pipe", "pipe"],
    });

    let settled = false;
    let stdout = "";

    const done = (result: Discovery | null) => {
      if (settled) return;
      settled = true;
      clearTimeout(timer);
      resolve(result);
    };

    const timer = setTimeout(() => {
      killProcessTree(proc);
      done(null);
    }, timeoutMs);

    proc.stdout?.on("data", (d) => (stdout += d.toString()));
    proc.on("error", () => done(null));
    proc.on("close", (code) => {
      const version = stdout.trim();
      if (code === 0 && version && (parseMajorVersion(version) ?? 0) >= 7) {
        done({ exe: "pwsh", shell: "pwsh", version });
      } else {
        done(null);
      }
    });
  });
}

async function findPowerShell(): Promise<Discovery> {
  if (cachedDiscovery) return cachedDiscovery;

  const discovery = await probePwsh7();
  if (!discovery) {
    throw new Error(
      "PowerShell 7+ is required, but `pwsh` was not found or is older than version 7. Install PowerShell 7 and ensure `pwsh` is on PATH."
    );
  }

  cachedDiscovery = discovery;
  return cachedDiscovery;
}

async function executePowerShell(options: {
  command: string;
  cwd: string;
  signal?: AbortSignal;
  timeoutMs?: number;
  maxOutputBytes?: number;
  maxLines?: number;
  onUpdate?: (partialResult: {
    content: Array<{ type: "text"; text: string }>;
    details: Partial<ExecResult> & { elapsedMs: number };
  }) => void;
}): Promise<ExecResult> {
  const discovery = await findPowerShell();
  const timeoutMs = clampInteger(options.timeoutMs, 120_000, 60 * 60 * 1000);
  const maxOutputBytes = clampInteger(options.maxOutputBytes, DEFAULT_STDOUT_MAX_BYTES, 500_000);
  const maxLines = clampInteger(options.maxLines, DEFAULT_STDOUT_MAX_LINES, 20_000);

  let cwdStat;
  try {
    cwdStat = await stat(options.cwd);
  } catch {
    throw new Error(`Working directory does not exist: ${options.cwd}`);
  }
  if (!cwdStat.isDirectory()) {
    throw new Error(`Working directory is not a directory: ${options.cwd}`);
  }

  const tempDir = await mkdtemp(join(tmpdir(), "pi-powershell-"));
  const args = await preparePowerShellArgs(options.command, tempDir);
  const stdoutCapture = new OutputCapture(join(tempDir, "stdout.txt"), maxOutputBytes, maxLines);
  const stderrCapture = new OutputCapture(join(tempDir, "stderr.txt"), STDERR_MAX_BYTES, STDERR_MAX_LINES);

  return new Promise((resolve) => {
    const proc = spawn(discovery.exe, args, {
      cwd: options.cwd,
      windowsHide: true,
      stdio: ["ignore", "pipe", "pipe"],
    });

    let settled = false;
    let timedOut = false;
    let aborted = false;
    let terminating = false;
    let forceTimer: NodeJS.Timeout | undefined;
    const startedAt = Date.now();
    const progressTimer = setInterval(() => {
      const elapsedMs = Date.now() - startedAt;
      const elapsedSeconds = Math.max(1, Math.floor(elapsedMs / 1000));
      options.onUpdate?.({
        content: [{ type: "text", text: `Running PowerShell... ${elapsedSeconds}s` }],
        details: {
          shell: discovery.shell,
          shellVersion: discovery.version,
          shellExe: discovery.exe,
          cwd: options.cwd,
          command: options.command,
          elapsedMs,
        },
      });
    }, 1000);

    const finalize = async (exitCode: number) => {
      if (settled) return;
      settled = true;
      clearTimeout(timer);
      clearInterval(progressTimer);
      if (forceTimer) clearTimeout(forceTimer);
      options.signal?.removeEventListener("abort", abortHandler);

      const out = await stdoutCapture.finish();
      const err = await stderrCapture.finish();
      const truncated = out.truncated || err.truncated;

      if (!truncated) {
        await rm(tempDir, { recursive: true, force: true }).catch(() => undefined);
      }

      resolve({
        shell: discovery.shell,
        shellVersion: discovery.version,
        shellExe: discovery.exe,
        cwd: options.cwd,
        command: options.command,
        exitCode,
        stdout: out.text,
        stderr: err.text,
        timedOut,
        aborted,
        truncated,
        stdoutFullPath: out.fullPath,
        stderrFullPath: err.fullPath,
      });
    };

    const terminate = (reason: "timeout" | "abort") => {
      if (settled) return;
      if (reason === "timeout") timedOut = true;
      if (reason === "abort") aborted = true;
      if (terminating) return;

      terminating = true;
      clearTimeout(timer);
      killProcessTree(proc);
      forceTimer = setTimeout(() => finalize(-1), 5_000);
    };

    const timer = setTimeout(() => terminate("timeout"), timeoutMs);
    const abortHandler = () => terminate("abort");

    if (options.signal?.aborted) {
      terminate("abort");
    } else {
      options.signal?.addEventListener("abort", abortHandler, { once: true });
    }

    proc.stdout?.on("data", (d) => stdoutCapture.append(d));
    proc.stderr?.on("data", (d) => stderrCapture.append(d));

    proc.on("error", (error) => {
      stderrCapture.append(error.message);
      finalize(-1);
    });

    proc.on("close", (code) => finalize(code ?? -1));
  });
}

export default function (pi: ExtensionAPI) {
  pi.registerTool({
    name: "powershell",
    label: "PowerShell",
    description:
      "Run a foreground, non-interactive PowerShell 7+ command on Windows via pwsh. Output is truncated to the last 500 lines or 20KB by default and full output is saved to a temp file when truncated.",
    promptSnippet:
      "Run a non-interactive PowerShell 7+ command on Windows via pwsh (supports cwd, timeout, abort, tail-truncated output)",
    promptGuidelines: [
      "Use powershell for Windows-native commands, C:\\ paths, $env variables, .ps1 execution, and Windows system inspection.",
      "The powershell tool requires PowerShell 7+ (`pwsh`) and does not fall back to Windows PowerShell 5.1.",
      "Do not use powershell for git, Unix-style text processing, or normal file reads/edits when bash/read/edit/write are more appropriate.",
      "Do not use powershell for interactive programs; it runs with -NonInteractive in the foreground.",
    ],
    parameters: Type.Object({
      command: Type.String({ description: "PowerShell command to execute" }),
      cwd: Type.Optional(
        Type.String({
          description:
            "Working directory (default: current working directory; relative paths resolve against it)",
        })
      ),
      timeoutMs: Type.Optional(
        Type.Integer({
          minimum: 1,
          maximum: 60 * 60 * 1000,
          description: "Timeout in milliseconds (default 120000, max 3600000)",
        })
      ),
      maxOutputBytes: Type.Optional(
        Type.Integer({
          minimum: 1,
          maximum: 500_000,
          description: "Maximum stdout bytes returned to the model (default 20000, max 500000)",
        })
      ),
      maxLines: Type.Optional(
        Type.Integer({
          minimum: 1,
          maximum: 20_000,
          description: "Maximum stdout lines returned to the model (default 500, max 20000)",
        })
      ),
    }, { additionalProperties: false }),
    renderCall(args, theme, context) {
      const text =
        context.lastComponent instanceof WidthAwareText
          ? context.lastComponent
          : new WidthAwareText();
      const command = typeof args.command === "string" ? args.command : "";
      const expanded = context.expanded ?? false;
      const preview = formatCommandPreview(command, expanded) || "...";
      const cwd = typeof args.cwd === "string" && args.cwd ? ` in ${args.cwd}` : "";

      text.setRenderText((width) => {
        if (expanded) {
          const title = theme.fg("toolTitle", theme.bold(`PS> ${command}`));
          const workingDirectory = cwd ? theme.fg("muted", cwd) : "";
          return `${title}${workingDirectory}\n`;
        }

        const rawContent = `PS> ${preview}${cwd}`;
        const suffixWidth = visibleWidth(`PS> ${cwd}`);
        if (suffixWidth >= width) {
          const compact = stripTerminalSequences(
            truncateToWidth(rawContent, Math.max(1, width), "…"),
          );
          return `${theme.fg("toolTitle", theme.bold(compact))}\n`;
        }

        const visiblePreview = stripTerminalSequences(
          truncateToWidth(preview, Math.max(1, width - suffixWidth), "…"),
        );
        const title = theme.fg("toolTitle", theme.bold(`PS> ${visiblePreview}`));
        const workingDirectory = cwd ? theme.fg("muted", cwd) : "";
        return `${title}${workingDirectory}\n`;
      });
      return text;
    },
    renderResult(result, options, theme, context) {
      const text =
        context.lastComponent instanceof WidthAwareText
          ? context.lastComponent
          : new WidthAwareText();
      const output = getTextContent(result);

      if (options.isPartial) {
        text.setRenderText(() => theme.fg("warning", output || "Running PowerShell..."));
        return text;
      }

      const summary = theme.fg(
        context.isError ? "error" : "muted",
        formatResultSummary(result.details as Partial<ExecResult> | undefined),
      );
      text.setRenderText((width) => {
        if (!output) return summary;
        if (options.expanded) {
          const styledOutput = output
            .split(/\r?\n/)
            .map((line) => (line ? theme.fg("toolOutput", line) : ""))
            .join("\n");
          return `${summary}\n${styledOutput}`;
        }

        const wrappedLines = wrapTextWithAnsi(output, Math.max(1, width));
        const previewLines = wrappedLines.slice(0, COLLAPSED_OUTPUT_MAX_LINES);
        const skippedLines = wrappedLines.length - previewLines.length;
        const styledPreviewLines = previewLines.map((line) =>
          line ? theme.fg("toolOutput", line) : "",
        );
        if (skippedLines > 0) {
          styledPreviewLines.push(
            theme.fg(
              "muted",
              `… (${skippedLines} earlier visual line${skippedLines === 1 ? "" : "s"}, ${keyText(
                "app.tools.expand",
              )} to expand)`,
            ),
          );
        }

        return `${summary}\n${styledPreviewLines.join("\n")}`;
      });
      return text;
    },
    async execute(_toolCallId, params, signal, onUpdate, ctx) {
      const result = await executePowerShell({
        command: params.command,
        cwd: normalizeCwd(params.cwd, ctx.cwd),
        signal,
        timeoutMs: params.timeoutMs,
        maxOutputBytes: params.maxOutputBytes,
        maxLines: params.maxLines,
        onUpdate,
      });

      const outputText = formatOutputForContent(result, "");

      if (result.aborted) {
        throw new Error(appendStatus(outputText, "Command aborted"));
      }
      if (result.timedOut) {
        throw new Error(appendStatus(outputText, `Command timed out after ${params.timeoutMs ?? 120_000}ms`));
      }
      if (result.exitCode !== 0) {
        throw new Error(appendStatus(outputText, `Command exited with code ${result.exitCode}`));
      }

      return {
        content: [{ type: "text", text: formatOutputForContent(result) }],
        details: result,
      };
    },
  });
}