import { NextRequest, NextResponse } from "next/server";
import fs from "fs";
import path from "path";
import { Client } from "@gradio/client";
import { query } from "../../../db";

export const maxDuration = 120; // Allow up to 120 seconds for slow model inference

export async function GET(req: NextRequest) {
  try {
    const { searchParams } = new URL(req.url);
    const action = searchParams.get("action") || "list";
    const runId = searchParams.get("runId");
    const imageType = searchParams.get("imageType"); // 'do', 'product' or null for all

    if (runId) {
      const runRes = await query(`
        SELECT id, image_path, engine, status, ocr_result, time_elapsed_ms, image_type, created_at
        FROM arena_runs
        WHERE id = $1
      `, [parseInt(runId)]);
      
      if (runRes.rowCount === 0) {
        return NextResponse.json({ success: false, error: "Run not found" }, { status: 404 });
      }
      return NextResponse.json({ success: true, run: runRes.rows[0] });
    }

    if (action === "stats") {
      let queryText = `
        SELECT 
          engine,
          COUNT(*)::integer as total_runs,
          COUNT(CASE WHEN status = 'done' THEN 1 END)::integer as success_runs,
          COUNT(CASE WHEN status = 'failed' THEN 1 END)::integer as failed_runs,
          ROUND(AVG(CASE WHEN status = 'done' THEN time_elapsed_ms END))::integer as avg_time_ms,
          MIN(CASE WHEN status = 'done' THEN time_elapsed_ms END)::integer as min_time_ms,
          MAX(CASE WHEN status = 'done' THEN time_elapsed_ms END)::integer as max_time_ms
        FROM arena_runs
      `;
      const params: any[] = [];
      if (imageType === "do" || imageType === "product") {
        queryText += ` WHERE image_type = $1`;
        params.push(imageType);
      }
      queryText += ` GROUP BY engine`;

      const statsRes = await query(queryText, params);
      return NextResponse.json({ success: true, stats: statsRes.rows });
    }

    const limit = parseInt(searchParams.get("limit") || "50");
    let queryText = `
      SELECT id, image_path, engine, status, time_elapsed_ms, image_type, created_at
      FROM arena_runs
    `;
    const params: any[] = [];
    if (imageType === "do" || imageType === "product") {
      queryText += ` WHERE image_type = $1`;
      params.push(imageType);
    }
    queryText += ` ORDER BY created_at DESC LIMIT $${params.length + 1}`;
    params.push(limit);

    const runsRes = await query(queryText, params);
    return NextResponse.json({ success: true, runs: runsRes.rows });
  } catch (error: any) {
    console.error("Failed to fetch arena runs/stats:", error);
    return NextResponse.json({ success: false, error: error.message }, { status: 500 });
  }
}

export async function POST(req: NextRequest) {
  const startTime = Date.now();
  let engine: string | undefined;
  let image: string | undefined;
  let imageType = "do";
  try {
    const body = await req.json().catch(() => ({}));
    engine = body.engine;
    image = body.image;

    if (!engine || !image) {
      return NextResponse.json({ error: "Missing engine or image" }, { status: 400 });
    }

    imageType = body.imageType || "do";
    if (typeof image === "string") {
      if (image.startsWith("/produk-pfm/") || image.includes("produk-pfm") || image.includes("Product")) {
        imageType = "product";
      } else if (image.startsWith("/do-pfm/") || image.includes("do-pfm")) {
        imageType = "do";
      }
    }

    let imageBuffer: Buffer;
    let base64Image = "";

    // 1. Resolve image (local file or base64)
    if (typeof image === "string" && (image.startsWith("/do-pfm/") || image.startsWith("/produk-pfm/"))) {
      // Resolve path in public folder
      const cleanPath = image.startsWith("/") ? image.slice(1) : image;
      const filePath = path.join(process.cwd(), "public", cleanPath);
      
      if (!fs.existsSync(filePath)) {
        return NextResponse.json({ error: `File not found on server: ${image}` }, { status: 404 });
      }
      imageBuffer = fs.readFileSync(filePath);
      base64Image = `data:image/jpeg;base64,${imageBuffer.toString("base64")}`;
    } else if (typeof image === "string" && image.startsWith("data:")) {
      // Base64 data URI
      base64Image = image;
      const base64Data = image.split(",")[1];
      imageBuffer = Buffer.from(base64Data, "base64");
    } else if (typeof image === "string") {
      // Raw base64 string
      base64Image = `data:image/jpeg;base64,${image}`;
      imageBuffer = Buffer.from(image, "base64");
    } else {
      return NextResponse.json({ error: "Invalid image format" }, { status: 400 });
    }

    let outputText = "";

    // 2. Route to the requested OCR engine
    if (engine === "deepseek") {
      const blob = new Blob([new Uint8Array(imageBuffer)], { type: "image/jpeg" });
      const gradioUrl = process.env.DEEPSEEK_GRADIO_URL || "http://host.docker.internal:7873/v2/";
      const client = await Client.connect(gradioUrl);
      const result = await client.predict(2, [blob, "Default", "Markdown", ""]);
      const data = result.data as any[];
      outputText = data[1] || data[0] || "";

    } else if (engine === "lightonocr") {
      const url = process.env.LIGHTONOCR_API_URL || "http://host.docker.internal:7678/layout-parsing";
      const res = await fetch(url, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          file: base64Image,
          useLayoutDetection: false
        })
      });
      if (!res.ok) {
        throw new Error(`LightOnOCR backend error: ${res.status} ${await res.text()}`);
      }
      const data = await res.json();
      outputText = data.result?.layoutParsingResults?.[0]?.markdown?.text || "";

    } else if (engine === "nemotron") {
      const url = process.env.NEMOTRON_API_URL || "http://host.docker.internal:8009/layout-parsing";
      const res = await fetch(url, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          file: base64Image,
          model: "Multilingual (en, zh, ja, ko, ru, …)",
          merge_level: "layout"
        })
      });
      if (!res.ok) {
        throw new Error(`Nemotron backend error: ${res.status} ${await res.text()}`);
      }
      const data = await res.json();
      outputText = data.result?.layoutParsingResults?.[0]?.markdown?.text || "";

    } else if (engine === "paddle") {
      const url = process.env.PIPELINE_URL || "http://paddleocr-pipeline-api:8090/layout-parsing";
      const rawB64 = base64Image.includes(",") ? base64Image.split(",")[1] : base64Image;
      const res = await fetch(url, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          file: rawB64,
          matchHistoryJob: false,
          useLayoutDetection: true,
          fileType: 1,
          useDocUnwarping: false,
          useDocOrientationClassify: false
        })
      });
      if (!res.ok) {
        throw new Error(`PaddleOCR backend error: ${res.status} ${await res.text()}`);
      }
      const data = await res.json();
      const pipelineResult = data.result || data;
      outputText = pipelineResult?.layoutParsingResults?.[0]?.markdown?.text || "";

    } else if (engine === "dots") {
      // Calling python API directly
      const url = process.env.DOTS_API_URL || "http://host.docker.internal:7872/layout-parsing";
      const res = await fetch(url, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          file: base64Image,
          promptLabel: "ocr",
          useLayoutDetection: true
        })
      });
      if (!res.ok) {
        throw new Error(`Dots OCR backend error: ${res.status} ${await res.text()}`);
      }
      const data = await res.json();
      outputText = data.result?.layoutParsingResults?.[0]?.markdown?.text || "";

    } else if (engine === "glm") {
      const gradioUrl = process.env.GLM_GRADIO_URL || "http://host.docker.internal:7875/";
      const client = await Client.connect(gradioUrl);
      const result = await client.predict(2, ["Text", base64Image, 1024, 60]);
      const data = result.data as any[];
      outputText = data[0] || "";

    } else {
      return NextResponse.json({ error: `Unknown engine: ${engine}` }, { status: 400 });
    }

    const elapsedMs = Date.now() - startTime;
    
    // Record successful run
    try {
      const loggedImagePath = (typeof image === "string" && image.startsWith("data:")) 
        ? `[Base64 Upload: ${image.length} chars]` 
        : (typeof image === "string" && image.length > 500) 
          ? `[Raw Base64: ${image.length} chars]` 
          : image;
      await query(
        `INSERT INTO arena_runs (image_path, engine, status, ocr_result, time_elapsed_ms, image_type)
         VALUES ($1, $2, $3, $4, $5, $6)`,
        [loggedImagePath, engine, "done", outputText, elapsedMs, imageType]
      );
    } catch (dbErr) {
      console.error("Failed to log success to arena_runs:", dbErr);
    }

    return NextResponse.json({
      success: true,
      text: outputText,
      elapsedMs
    });

  } catch (error: any) {
    console.error("OCR Arena proxy error:", error);
    const elapsedMs = Date.now() - startTime;
    
    // Record failed run
    try {
      const loggedImagePath = (typeof image === "string" && image.startsWith("data:")) 
        ? `[Base64 Upload: ${image.length} chars]` 
        : (typeof image === "string" && image.length > 500) 
          ? `[Raw Base64: ${image.length} chars]` 
          : image;
      await query(
        `INSERT INTO arena_runs (image_path, engine, status, ocr_result, time_elapsed_ms, image_type)
         VALUES ($1, $2, $3, $4, $5, $6)`,
        [loggedImagePath || "unknown", engine || "unknown", "failed", error.message || "Unknown error", elapsedMs, imageType]
      );
    } catch (dbErr) {
      console.error("Failed to log failure to arena_runs:", dbErr);
    }

    return NextResponse.json({
      success: false,
      error: error.message || "Failed to process OCR request"
    }, { status: 500 });
  }
}