import { spawn } from "node:child_process";
import { createHash } from "node:crypto";
import { access, mkdir, readdir, rm, stat } from "node:fs/promises";
import { constants } from "node:fs";
import { join, resolve } from "node:path";

/**
 * Bridge to the Grok Build CLI for reference-image generation.
 *
 * `/imagine` is not a CLI subcommand — it is a bundled *skill* that documents
 * two agent tools, `image_gen` and `image_edit`. Neither is reachable from the
 * shell directly, so the bridge drives a headless single-turn session
 * (`grok -p … --output-format json`) and lets the agent make the tool call,
 * then recovers the file it wrote.
 *
 * Two consequences shape the design:
 *
 *  1. The agent decides where the file lands. We give it a dedicated working
 *     directory and diff the contents rather than trusting the path it reports
 *     in prose, which is free-form and occasionally wrapped or annotated.
 *  2. Generation costs real money and takes ~60-90s, so results are cached by
 *     prompt hash. Re-running the pipeline on the same prompt is free.
 */

export interface GrokBridgeOptions {
  /** Executable name or absolute path. */
  command?: string;
  /** Where generated images are written and cached. */
  outputDir?: string;
  timeoutMs?: number;
  model?: string;
}

export interface ImageRequest {
  prompt: string;
  aspectRatio?: "1:1" | "16:9" | "9:16" | "4:3" | "3:4" | "auto";
  /** Existing image(s) to edit. Present means `image_edit` rather than `image_gen`. */
  sourceImages?: string[];
  /**
   * Re-run even if this prompt is already cached, and replace the cache entry
   * with the new image. Costs money every time — see {@link GrokBridge.generate}.
   *
   * Deliberately absent from {@link hashRequest}: a forced run has to land in
   * the *same* cache slot it is replacing, or the stale image would keep
   * winning on the next unforced call.
   */
  force?: boolean;
}

export interface ImageResult {
  path: string;
  prompt: string;
  cached: boolean;
  costUsd: number;
  elapsedMs: number;
}

export class GrokUnavailableError extends Error {
  constructor(command: string) {
    super(
      `Grok CLI not found (looked for "${command}"). Install it, or pass an explicit image ` +
        `with --image to skip generation.`,
    );
    this.name = "GrokUnavailableError";
  }
}

export class GrokGenerationError extends Error {
  constructor(
    message: string,
    readonly stderr: string,
  ) {
    super(message);
    this.name = "GrokGenerationError";
  }
}

interface GrokJsonResponse {
  text?: string;
  stopReason?: string;
  total_cost_usd?: number;
  sessionId?: string;
}

const IMAGE_EXTENSIONS = new Set([".jpg", ".jpeg", ".png", ".webp"]);

export class GrokBridge {
  private readonly command: string;
  private readonly outputDir: string;
  private readonly timeoutMs: number;
  private readonly model: string | undefined;

  constructor(options: GrokBridgeOptions = {}) {
    this.command = options.command ?? process.env["GROK_BIN"] ?? "grok";
    this.outputDir = resolve(options.outputDir ?? "assets/references");
    this.timeoutMs = options.timeoutMs ?? 5 * 60_000;
    this.model = options.model;
  }

  /** True when the CLI is on PATH and responds. */
  async isAvailable(): Promise<boolean> {
    try {
      await this.exec(["--version"], 20_000);
      return true;
    } catch {
      return false;
    }
  }

  /**
   * Generate (or edit) a reference image.
   *
   * The prompt is wrapped in an instruction that pins the agent to a single
   * tool call and a predictable output location — left to its own devices it
   * may narrate, ask questions, or save somewhere else.
   *
   * `request.force` skips the cache and overwrites it. Worth reaching for
   * because the cache key is the prompt, not the mesh: edit a character's face
   * or physique without renaming it and the prompt is unchanged, so an unforced
   * run happily returns the texture painted for the *old* body.
   */
  async generate(request: ImageRequest): Promise<ImageResult> {
    const started = Date.now();
    const key = hashRequest(request);
    const workDir = join(this.outputDir, key);

    if (request.force) {
      // Clear the slot first, so the fresh image is unambiguously the one that
      // `findCached` picks up next time (it takes the first name in sort order,
      // which a leftover file could otherwise win).
      await this.clearCache(workDir);
    } else {
      const cachedPath = await this.findCached(workDir);
      if (cachedPath) {
        return { path: cachedPath, prompt: request.prompt, cached: true, costUsd: 0, elapsedMs: 0 };
      }
    }

    await mkdir(workDir, { recursive: true });
    const before = await listImages(workDir);

    const args = ["-p", this.buildInstruction(request), "--output-format", "json", "--always-approve", "--cwd", workDir];
    if (this.model) args.push("--model", this.model);

    const { stdout, stderr } = await this.exec(args, this.timeoutMs);

    let response: GrokJsonResponse = {};
    try {
      response = JSON.parse(stdout) as GrokJsonResponse;
    } catch {
      throw new GrokGenerationError("Grok did not return parseable JSON", stderr || stdout.slice(0, 500));
    }

    // Prefer a file that actually appeared over the path quoted in prose.
    const after = await listImages(workDir);
    const created = after.filter((file) => !before.includes(file));
    const path = created[0] ?? extractPath(response.text ?? "");

    if (!path) {
      throw new GrokGenerationError(
        `Grok produced no image. It said: ${(response.text ?? "").slice(0, 300)}`,
        stderr,
      );
    }

    return {
      path: created[0] ? join(workDir, created[0]) : path,
      prompt: request.prompt,
      cached: false,
      costUsd: response.total_cost_usd ?? 0,
      elapsedMs: Date.now() - started,
    };
  }

  private buildInstruction(request: ImageRequest): string {
    const aspect = request.aspectRatio ?? "3:4";

    if (request.sourceImages?.length) {
      return [
        `Call image_edit exactly once with these source images: ${request.sourceImages.join(", ")}.`,
        `Transformation: ${request.prompt}`,
        `Save the resulting image into the current working directory.`,
        `Do not ask questions. Reply with only the absolute file path on the final line.`,
      ].join("\n");
    }

    return [
      `Call image_gen exactly once with aspect_ratio ${aspect}.`,
      `Prompt: ${request.prompt}`,
      `Save the resulting image into the current working directory.`,
      `Do not ask questions and do not generate more than one image.`,
      `Reply with only the absolute file path on the final line.`,
    ].join("\n");
  }

  private async findCached(workDir: string): Promise<string | null> {
    try {
      const images = await listImages(workDir);
      return images[0] ? join(workDir, images[0]) : null;
    } catch {
      return null;
    }
  }

  /**
   * Drop the images in one cache slot.
   *
   * Only image files, and only inside the hashed directory — never the slot
   * itself or anything else the agent left behind. `workDir` is built from a
   * hex digest, so it cannot escape `outputDir`.
   */
  private async clearCache(workDir: string): Promise<void> {
    for (const name of await listImages(workDir)) {
      await rm(join(workDir, name), { force: true });
    }
  }

  private exec(args: string[], timeoutMs: number): Promise<{ stdout: string; stderr: string }> {
    return new Promise((resolvePromise, reject) => {
      let child: ReturnType<typeof spawn>;
      try {
        child = spawn(this.command, args, { stdio: ["ignore", "pipe", "pipe"] });
      } catch {
        reject(new GrokUnavailableError(this.command));
        return;
      }

      let stdout = "";
      let stderr = "";
      let settled = false;

      const timer = setTimeout(() => {
        if (settled) return;
        settled = true;
        child.kill("SIGKILL");
        reject(new GrokGenerationError(`Grok timed out after ${timeoutMs}ms`, stderr));
      }, timeoutMs);

      child.stdout?.on("data", (chunk: Buffer) => (stdout += chunk.toString()));
      child.stderr?.on("data", (chunk: Buffer) => (stderr += chunk.toString()));

      child.on("error", (error: NodeJS.ErrnoException) => {
        if (settled) return;
        settled = true;
        clearTimeout(timer);
        reject(error.code === "ENOENT" ? new GrokUnavailableError(this.command) : error);
      });

      child.on("close", (code) => {
        if (settled) return;
        settled = true;
        clearTimeout(timer);
        if (code !== 0) {
          reject(new GrokGenerationError(`Grok exited with code ${code}`, stderr));
          return;
        }
        resolvePromise({ stdout, stderr });
      });
    });
  }
}

function hashRequest(request: ImageRequest): string {
  return createHash("sha256")
    .update(JSON.stringify([request.prompt, request.aspectRatio, request.sourceImages]))
    .digest("hex")
    .slice(0, 16);
}

async function listImages(dir: string): Promise<string[]> {
  const entries = await readdir(dir).catch(() => [] as string[]);
  const images: string[] = [];
  for (const entry of entries) {
    const dot = entry.lastIndexOf(".");
    if (dot === -1 || !IMAGE_EXTENSIONS.has(entry.slice(dot).toLowerCase())) continue;
    const info = await stat(join(dir, entry)).catch(() => null);
    if (info?.isFile()) images.push(entry);
  }
  return images.sort();
}

/** Last absolute-looking path in the agent's prose, as a fallback. */
function extractPath(text: string): string | null {
  const matches = text.match(/\/[^\s"'`]+\.(?:jpg|jpeg|png|webp)/gi);
  return matches?.[matches.length - 1] ?? null;
}

/** Confirms a path exists and is readable before the pipeline commits to it. */
export async function ensureReadable(path: string): Promise<void> {
  await access(path, constants.R_OK);
}
