import { spawn } from "node:child_process";
import { promises as fs } from "node:fs";
import path from "node:path";
import { type ChatMessage, type PromptPreprocessorController } from "@lmstudio/sdk";
import { configSchematics } from "./config";
import { globalConfigSchematics } from "./globalConfig";

const VIDEO_EXTENSIONS = new Set([
  ".3gp",
  ".asf",
  ".avi",
  ".f4v",
  ".flv",
  ".m2ts",
  ".m4v",
  ".mkv",
  ".mov",
  ".mp4",
  ".mpg",
  ".mpeg",
  ".mxf",
  ".ogv",
  ".rm",
  ".rmvb",
  ".ts",
  ".vob",
  ".webm",
  ".wmv",
  ".y4m",
]);

const RESOLUTION_TARGETS: Record<string, number> = {
  "300k": 300_000,
  "600k": 600_000,
  "1mpx": 1_000_000,
  "2mpx": 2_000_000,
};

interface VideoInfo {
  durationSec: number | null;
  width: number | null;
  height: number | null;
}

function isVideoFile(name: string): boolean {
  return VIDEO_EXTENSIONS.has(path.extname(name).toLowerCase());
}

function toFiniteNumber(value: unknown): number | null {
  if (typeof value === "number" && Number.isFinite(value)) {
    return value;
  }
  if (typeof value === "string") {
    const parsed = Number.parseFloat(value);
    if (Number.isFinite(parsed)) {
      return parsed;
    }
  }
  return null;
}

function runCommand(
  command: string,
  args: string[],
  signal?: AbortSignal,
): Promise<{ code: number | null; stdout: string; stderr: string }> {
  return new Promise((resolve, reject) => {
    const proc = spawn(command, args);
    let stdout = "";
    let stderr = "";
    const onAbort = () => {
      proc.kill("SIGTERM");
    };
    if (signal) {
      if (signal.aborted) {
        proc.kill("SIGTERM");
      } else {
        signal.addEventListener("abort", onAbort, { once: true });
      }
    }
    const cleanup = () => {
      if (signal) {
        signal.removeEventListener("abort", onAbort);
      }
    };
    proc.stdout.on("data", (chunk: Buffer) => {
      stdout += chunk.toString();
    });
    proc.stderr.on("data", (chunk: Buffer) => {
      stderr += chunk.toString();
    });
    proc.on("error", (err: Error) => {
      cleanup();
      reject(new Error(`Failed to run "${command}": ${err.message}`));
    });
    proc.on("close", code => {
      cleanup();
      resolve({ code, stdout, stderr });
    });
  });
}

async function probeVideo(
  ffprobePath: string,
  videoPath: string,
  signal?: AbortSignal,
): Promise<VideoInfo> {
  const args = [
    "-v",
    "error",
    "-select_streams",
    "v:0",
    "-show_entries",
    "stream=width,height,duration",
    "-show_entries",
    "format=duration",
    "-of",
    "json",
    videoPath,
  ];
  const { code, stdout, stderr } = await runCommand(ffprobePath, args, signal);
  if (code !== 0) {
    throw new Error(`ffprobe failed with exit code ${code}. ${(stderr.trim() || stdout.trim()).slice(0, 500)}`.trim());
  }
  let parsed: unknown;
  try {
    parsed = JSON.parse(stdout);
  } catch {
    throw new Error("Could not parse ffprobe output.");
  }
  const obj = parsed as { streams?: Array<Record<string, unknown>>; format?: Record<string, unknown> };
  const stream = obj.streams?.[0];
  if (!stream) {
    throw new Error(`No video stream found in "${path.basename(videoPath)}".`);
  }
  return {
    durationSec: toFiniteNumber(stream.duration) ?? toFiniteNumber(obj.format?.duration),
    width: toFiniteNumber(stream.width),
    height: toFiniteNumber(stream.height),
  };
}

function computeScale(
  width: number | null,
  height: number | null,
  resolutionKey: string,
): { w: number; h: number } | null {
  const target = RESOLUTION_TARGETS[resolutionKey];
  if (!target || !width || !height) {
    return null;
  }
  const currentPixels = width * height;
  if (currentPixels <= target) {
    return null;
  }
  const factor = Math.sqrt(target / currentPixels);
  const w = Math.max(2, Math.round((width * factor) / 2) * 2);
  const h = Math.max(2, Math.round((height * factor) / 2) * 2);
  return { w, h };
}

function sanitizeDirName(name: string): string {
  const cleaned = path.basename(name).replace(/[^a-zA-Z0-9._-]/g, "_").slice(0, 64);
  return cleaned || "video";
}

function formatDuration(totalSeconds: number): string {
  const seconds = Math.round(totalSeconds);
  if (seconds < 60) {
    return `${seconds}s`;
  }
  const minutes = Math.floor(seconds / 60);
  const rest = seconds % 60;
  if (minutes < 60) {
    return rest > 0 ? `${minutes}m ${rest}s` : `${minutes}m`;
  }
  const hours = Math.floor(minutes / 60);
  const minRest = minutes % 60;
  return minRest > 0 ? `${hours}h ${minRest}m` : `${hours}h`;
}

export async function preprocess(ctl: PromptPreprocessorController, userMessage: ChatMessage) {
  const allFiles = userMessage.getFiles(ctl.client);
  if (!allFiles.some(f => isVideoFile(f.name))) {
    return userMessage;
  }

  const config = ctl.getPluginConfig(configSchematics);
  const globalConfig = ctl.getGlobalPluginConfig(globalConfigSchematics);
  const ffmpegPath = (globalConfig.get("ffmpegPath") || "ffmpeg").trim();
  const ffprobePath = (globalConfig.get("ffprobePath") || "ffprobe").trim();

  const samplingMode = config.get("samplingMode");
  const fpsValue = config.get("fpsValue");
  const frameCount = config.get("frameCount");
  const resolution = config.get("resolution");
  const maxFramesPerVideo = Math.max(1, Math.round(config.get("maxFramesPerVideo")));

  const originalText = userMessage.getText();
  const videoFiles = userMessage.consumeFiles(ctl.client, f => isVideoFile(f.name));
  const notes: string[] = [];

  for (const video of videoFiles) {
    ctl.guardAbort();
    const status = ctl.createStatus({ status: "loading", text: `Converting ${video.name} to image frames...` });
    let workDir: string | null = null;
    try {
      status.setState({ status: "loading", text: `Analyzing ${video.name} with ffprobe...` });
      const videoPath = await video.getFilePath();
      const info = await probeVideo(ffprobePath, videoPath, ctl.abortSignal);

      let effectiveFps: number;
      if (samplingMode === "frameCount") {
        if (!info.durationSec || info.durationSec <= 0) {
          status.setState({
            status: "loading",
            text: `${video.name}: unknown duration, falling back to fixed FPS (${fpsValue} fps).`,
          });
          effectiveFps = fpsValue;
        } else {
          const requested = Math.max(1, Math.round(frameCount));
          if (requested > maxFramesPerVideo) {
            status.setState({
              status: "loading",
              text: `${video.name}: frame count ${requested} exceeds the cap of ${maxFramesPerVideo}, using ${maxFramesPerVideo}.`,
            });
          }
          const target = Math.min(requested, maxFramesPerVideo);
          effectiveFps = (target - 1) / info.durationSec;
        }
      } else if (info.durationSec && info.durationSec > 0) {
        const estimated = info.durationSec * fpsValue;
        if (estimated > maxFramesPerVideo) {
          status.setState({
            status: "loading",
            text: `${video.name}: ${fpsValue} fps would produce ~${Math.round(estimated)} frames, capping at ${maxFramesPerVideo}.`,
          });
          effectiveFps = maxFramesPerVideo / info.durationSec;
        } else {
          effectiveFps = fpsValue;
        }
      } else {
        effectiveFps = fpsValue;
      }

      const scale = computeScale(info.width, info.height, resolution);

      workDir = path.join(ctl.getWorkingDirectory(), `video-frames-${sanitizeDirName(video.name)}`);
      await fs.mkdir(workDir, { recursive: true });

      status.setState({ status: "loading", text: `Extracting frames from ${video.name} with ffmpeg...` });
      const filters = [`fps=${effectiveFps}`];
      if (scale) {
        filters.push(`scale=${scale.w}:${scale.h}`);
      }
      const outPattern = path.join(workDir, "frame_%04d.jpg");
      const args = [
        "-hide_banner",
        "-loglevel",
        "error",
        "-y",
        "-i",
        videoPath,
        "-vf",
        filters.join(","),
        "-q:v",
        "2",
        outPattern,
      ];
      const { code, stderr } = await runCommand(ffmpegPath, args, ctl.abortSignal);
      if (code !== 0) {
        throw new Error(`ffmpeg failed with exit code ${code}. ${(stderr.trim() || "").slice(0, 500)}`.trim());
      }

      let frameNames = (await fs.readdir(workDir)).filter(f => /^frame_\d{4}\.jpg$/.test(f)).sort();
      if (frameNames.length === 0) {
        throw new Error("ffmpeg did not produce any frames.");
      }
      if (frameNames.length > maxFramesPerVideo) {
        const step = frameNames.length / maxFramesPerVideo;
        frameNames = Array.from({ length: maxFramesPerVideo }, (_, i) => frameNames[Math.floor(i * step)]);
      }

      status.setState({ status: "loading", text: `Uploading ${frameNames.length} frame(s) of ${video.name}...` });
      let uploaded = 0;
      for (const frameName of frameNames) {
        ctl.guardAbort();
        const handle = await ctl.client.files.prepareImage(path.join(workDir, frameName));
        userMessage.appendFile(handle);
        uploaded += 1;
        status.setState({
          status: "loading",
          text: `Uploading frames of ${video.name} (${uploaded}/${frameNames.length})...`,
        });
      }

      const header = [`Video "${video.name}"`];
      if (info.durationSec) {
        header.push(`(${formatDuration(info.durationSec)})`);
      }
      notes.push(
        `${header.join(" ")} was converted to ${frameNames.length} image frame(s) in chronological order${
          scale ? `, downscaled to ${scale.w}x${scale.h}` : ""
        }. Treat them as a sequence of stills taken from the video.`,
      );
      status.setState({
        status: "done",
        text: `${video.name}: extracted ${frameNames.length} frame(s)${scale ? ` at ${scale.w}x${scale.h}` : ""}.`,
      });
    } catch (err) {
      const message = err instanceof Error ? err.message : String(err);
      status.setState({ status: "error", text: `${video.name}: ${message}` });
      throw new Error(`Failed to process video "${video.name}": ${message}`);
    } finally {
      if (workDir) {
        await fs.rm(workDir, { recursive: true, force: true }).catch(() => {});
      }
    }
  }

  const finalText = notes.length > 0 ? `${notes.join("\n\n")}\n\n${originalText}` : originalText;
  userMessage.replaceText(finalText);
  return userMessage;
}

import { spawn } from "node:child_process";
import { promises as fs } from "node:fs";
import path from "node:path";
import { type ChatMessage, type PromptPreprocessorController } from "@lmstudio/sdk";
import { configSchematics } from "./config";
import { globalConfigSchematics } from "./globalConfig";

const VIDEO_EXTENSIONS = new Set([
  ".3gp",
  ".asf",
  ".avi",
  ".f4v",
  ".flv",
  ".m2ts",
  ".m4v",
  ".mkv",
  ".mov",
  ".mp4",
  ".mpg",
  ".mpeg",
  ".mxf",
  ".ogv",
  ".rm",
  ".rmvb",
  ".ts",
  ".vob",
  ".webm",
  ".wmv",
  ".y4m",
]);

const RESOLUTION_TARGETS: Record<string, number> = {
  "300k": 300_000,
  "600k": 600_000,
  "1mpx": 1_000_000,
  "2mpx": 2_000_000,
};

interface VideoInfo {
  durationSec: number | null;
  width: number | null;
  height: number | null;
}

function isVideoFile(name: string): boolean {
  return VIDEO_EXTENSIONS.has(path.extname(name).toLowerCase());
}

function toFiniteNumber(value: unknown): number | null {
  if (typeof value === "number" && Number.isFinite(value)) {
    return value;
  }
  if (typeof value === "string") {
    const parsed = Number.parseFloat(value);
    if (Number.isFinite(parsed)) {
      return parsed;
    }
  }
  return null;
}

function runCommand(
  command: string,
  args: string[],
  signal?: AbortSignal,
): Promise<{ code: number | null; stdout: string; stderr: string }> {
  return new Promise((resolve, reject) => {
    const proc = spawn(command, args);
    let stdout = "";
    let stderr = "";
    const onAbort = () => {
      proc.kill("SIGTERM");
    };
    if (signal) {
      if (signal.aborted) {
        proc.kill("SIGTERM");
      } else {
        signal.addEventListener("abort", onAbort, { once: true });
      }
    }
    const cleanup = () => {
      if (signal) {
        signal.removeEventListener("abort", onAbort);
      }
    };
    proc.stdout.on("data", (chunk: Buffer) => {
      stdout += chunk.toString();
    });
    proc.stderr.on("data", (chunk: Buffer) => {
      stderr += chunk.toString();
    });
    proc.on("error", (err: Error) => {
      cleanup();
      reject(new Error(`Failed to run "${command}": ${err.message}`));
    });
    proc.on("close", code => {
      cleanup();
      resolve({ code, stdout, stderr });
    });
  });
}

async function probeVideo(
  ffprobePath: string,
  videoPath: string,
  signal?: AbortSignal,
): Promise<VideoInfo> {
  const args = [
    "-v",
    "error",
    "-select_streams",
    "v:0",
    "-show_entries",
    "stream=width,height,duration",
    "-show_entries",
    "format=duration",
    "-of",
    "json",
    videoPath,
  ];
  const { code, stdout, stderr } = await runCommand(ffprobePath, args, signal);
  if (code !== 0) {
    throw new Error(`ffprobe failed with exit code ${code}. ${(stderr.trim() || stdout.trim()).slice(0, 500)}`.trim());
  }
  let parsed: unknown;
  try {
    parsed = JSON.parse(stdout);
  } catch {
    throw new Error("Could not parse ffprobe output.");
  }
  const obj = parsed as { streams?: Array<Record<string, unknown>>; format?: Record<string, unknown> };
  const stream = obj.streams?.[0];
  if (!stream) {
    throw new Error(`No video stream found in "${path.basename(videoPath)}".`);
  }
  return {
    durationSec: toFiniteNumber(stream.duration) ?? toFiniteNumber(obj.format?.duration),
    width: toFiniteNumber(stream.width),
    height: toFiniteNumber(stream.height),
  };
}

function computeScale(
  width: number | null,
  height: number | null,
  resolutionKey: string,
): { w: number; h: number } | null {
  const target = RESOLUTION_TARGETS[resolutionKey];
  if (!target || !width || !height) {
    return null;
  }
  const currentPixels = width * height;
  if (currentPixels <= target) {
    return null;
  }
  const factor = Math.sqrt(target / currentPixels);
  const w = Math.max(2, Math.round((width * factor) / 2) * 2);
  const h = Math.max(2, Math.round((height * factor) / 2) * 2);
  return { w, h };
}

function sanitizeDirName(name: string): string {
  const cleaned = path.basename(name).replace(/[^a-zA-Z0-9._-]/g, "_").slice(0, 64);
  return cleaned || "video";
}

function formatDuration(totalSeconds: number): string {
  const seconds = Math.round(totalSeconds);
  if (seconds < 60) {
    return `${seconds}s`;
  }
  const minutes = Math.floor(seconds / 60);
  const rest = seconds % 60;
  if (minutes < 60) {
    return rest > 0 ? `${minutes}m ${rest}s` : `${minutes}m`;
  }
  const hours = Math.floor(minutes / 60);
  const minRest = minutes % 60;
  return minRest > 0 ? `${hours}h ${minRest}m` : `${hours}h`;
}

export async function preprocess(ctl: PromptPreprocessorController, userMessage: ChatMessage) {
  const allFiles = userMessage.getFiles(ctl.client);
  if (!allFiles.some(f => isVideoFile(f.name))) {
    return userMessage;
  }

  const config = ctl.getPluginConfig(configSchematics);
  const globalConfig = ctl.getGlobalPluginConfig(globalConfigSchematics);
  const ffmpegPath = (globalConfig.get("ffmpegPath") || "ffmpeg").trim();
  const ffprobePath = (globalConfig.get("ffprobePath") || "ffprobe").trim();

  const samplingMode = config.get("samplingMode");
  const fpsValue = config.get("fpsValue");
  const frameCount = config.get("frameCount");
  const resolution = config.get("resolution");
  const maxFramesPerVideo = Math.max(1, Math.round(config.get("maxFramesPerVideo")));

  const originalText = userMessage.getText();
  const videoFiles = userMessage.consumeFiles(ctl.client, f => isVideoFile(f.name));
  const notes: string[] = [];

  for (const video of videoFiles) {
    ctl.guardAbort();
    const status = ctl.createStatus({ status: "loading", text: `Converting ${video.name} to image frames...` });
    let workDir: string | null = null;
    try {
      status.setState({ status: "loading", text: `Analyzing ${video.name} with ffprobe...` });
      const videoPath = await video.getFilePath();
      const info = await probeVideo(ffprobePath, videoPath, ctl.abortSignal);

      let effectiveFps: number;
      if (samplingMode === "frameCount") {
        if (!info.durationSec || info.durationSec <= 0) {
          status.setState({
            status: "loading",
            text: `${video.name}: unknown duration, falling back to fixed FPS (${fpsValue} fps).`,
          });
          effectiveFps = fpsValue;
        } else {
          const requested = Math.max(1, Math.round(frameCount));
          if (requested > maxFramesPerVideo) {
            status.setState({
              status: "loading",
              text: `${video.name}: frame count ${requested} exceeds the cap of ${maxFramesPerVideo}, using ${maxFramesPerVideo}.`,
            });
          }
          const target = Math.min(requested, maxFramesPerVideo);
          effectiveFps = (target - 1) / info.durationSec;
        }
      } else if (info.durationSec && info.durationSec > 0) {
        const estimated = info.durationSec * fpsValue;
        if (estimated > maxFramesPerVideo) {
          status.setState({
            status: "loading",
            text: `${video.name}: ${fpsValue} fps would produce ~${Math.round(estimated)} frames, capping at ${maxFramesPerVideo}.`,
          });
          effectiveFps = maxFramesPerVideo / info.durationSec;
        } else {
          effectiveFps = fpsValue;
        }
      } else {
        effectiveFps = fpsValue;
      }

      const scale = computeScale(info.width, info.height, resolution);

      workDir = path.join(ctl.getWorkingDirectory(), `video-frames-${sanitizeDirName(video.name)}`);
      await fs.mkdir(workDir, { recursive: true });

      status.setState({ status: "loading", text: `Extracting frames from ${video.name} with ffmpeg...` });
      const filters = [`fps=${effectiveFps}`];
      if (scale) {
        filters.push(`scale=${scale.w}:${scale.h}`);
      }
      const outPattern = path.join(workDir, "frame_%04d.jpg");
      const args = [
        "-hide_banner",
        "-loglevel",
        "error",
        "-y",
        "-i",
        videoPath,
        "-vf",
        filters.join(","),
        "-q:v",
        "2",
        outPattern,
      ];
      const { code, stderr } = await runCommand(ffmpegPath, args, ctl.abortSignal);
      if (code !== 0) {
        throw new Error(`ffmpeg failed with exit code ${code}. ${(stderr.trim() || "").slice(0, 500)}`.trim());
      }

      let frameNames = (await fs.readdir(workDir)).filter(f => /^frame_\d{4}\.jpg$/.test(f)).sort();
      if (frameNames.length === 0) {
        throw new Error("ffmpeg did not produce any frames.");
      }
      if (frameNames.length > maxFramesPerVideo) {
        const step = frameNames.length / maxFramesPerVideo;
        frameNames = Array.from({ length: maxFramesPerVideo }, (_, i) => frameNames[Math.floor(i * step)]);
      }

      status.setState({ status: "loading", text: `Uploading ${frameNames.length} frame(s) of ${video.name}...` });
      let uploaded = 0;
      for (const frameName of frameNames) {
        ctl.guardAbort();
        const handle = await ctl.client.files.prepareImage(path.join(workDir, frameName));
        userMessage.appendFile(handle);
        uploaded += 1;
        status.setState({
          status: "loading",
          text: `Uploading frames of ${video.name} (${uploaded}/${frameNames.length})...`,
        });
      }

      const header = [`Video "${video.name}"`];
      if (info.durationSec) {
        header.push(`(${formatDuration(info.durationSec)})`);
      }
      notes.push(
        `${header.join(" ")} was converted to ${frameNames.length} image frame(s) in chronological order${
          scale ? `, downscaled to ${scale.w}x${scale.h}` : ""
        }. Treat them as a sequence of stills taken from the video.`,
      );
      status.setState({
        status: "done",
        text: `${video.name}: extracted ${frameNames.length} frame(s)${scale ? ` at ${scale.w}x${scale.h}` : ""}.`,
      });
    } catch (err) {
      const message = err instanceof Error ? err.message : String(err);
      status.setState({ status: "error", text: `${video.name}: ${message}` });
      throw new Error(`Failed to process video "${video.name}": ${message}`);
    } finally {
      if (workDir) {
        await fs.rm(workDir, { recursive: true, force: true }).catch(() => {});
      }
    }
  }

  const finalText = notes.length > 0 ? `${notes.join("\n\n")}\n\n${originalText}` : originalText;
  userMessage.replaceText(finalText);
  return userMessage;
}