import { readFile } from "fs/promises";
import { basename } from "path";
import { inflateRawSync } from "zlib";
import { formatServiceError, isAbortError } from "../shared/errors";
import { fetchJson } from "../shared/http";
import { MINERU_URL } from "../types";

export interface MinerUErrorResponse {
  code: number;
  msg: string;
  trace_id: string;
  data?: Record<string, any>;
}

export interface MinerUSubmitResponse {
  code: number;
  msg: string;
  trace_id: string;
  data: { task_id: string };
}

export interface MinerUTaskResult {
  code: number;
  msg: string;
  trace_id: string;
  data: {
    task_id: string;
    state: "done" | "pending" | "running" | "failed" | "converting";
    err_msg?: string;
    full_zip_url?: string;
    extract_progress?: {
      extracted_pages: number;
      total_pages: number;
      start_time: string;
    };
  };
}

export interface MinerUBatchUrlResponse {
  code: number;
  msg: string;
  trace_id: string;
  data: { batch_id: string; file_urls: string[] };
}

export interface MinerUBatchResult {
  code: number;
  msg: string;
  trace_id: string;
  data: {
    batch_id: string;
    extract_result: Array<{
      file_name: string;
      state: "done" | "waiting-file" | "pending" | "running" | "failed" | "converting";
      err_msg?: string;
      full_zip_url?: string;
      extract_progress?: {
        extracted_pages: number;
        total_pages: number;
        start_time: string;
      };
    }>;
  };
}

const POLL_INTERVAL_MS = 3000;
const MAX_POLL_ATTEMPTS = 40; // 40 * 3s = 120s max

interface ExtractOptions {
  modelVersion?: "pipeline" | "vlm" | "MinerU-HTML";
  pageRanges?: string;
  enableFormula?: boolean;
  enableTable?: boolean;
  language?: string;
}

async function submitTask(
  url: string,
  apiKey: string,
  signal?: AbortSignal,
  options?: ExtractOptions,
): Promise<{ taskId: string; error?: string }> {
  try {
    const body: Record<string, any> = { url };
    if (options?.modelVersion) body.model_version = options.modelVersion;
    if (options?.pageRanges) body.page_ranges = options.pageRanges;
    if (options?.enableFormula !== undefined) body.enable_formula = options.enableFormula;
    if (options?.enableTable !== undefined) body.enable_table = options.enableTable;
    if (options?.language) body.language = options.language;

    const data = await fetchJson<MinerUSubmitResponse>(
      `${MINERU_URL}/api/v4/extract/task`,
      {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          "Authorization": `Bearer ${apiKey}`,
        },
        body: JSON.stringify(body),
        signal,
      },
    );

    if (data.code !== 0) {
      return { taskId: "", error: `MinerU error ${data.code}: ${data.msg}` };
    }

    return { taskId: data.data.task_id };
  } catch (error) {
    if (isAbortError(error)) throw error;
    return { taskId: "", error: formatServiceError("MinerU", error) };
  }
}

async function requestUploadUrl(
  fileName: string,
  apiKey: string,
  signal?: AbortSignal,
  options?: ExtractOptions,
): Promise<{ batchId: string; uploadUrl: string; error?: string }> {
  try {
    // page_ranges is a per-file field on the batch endpoint; the rest stay top-level.
    const file: Record<string, any> = { name: fileName };
    if (options?.pageRanges) file.page_ranges = options.pageRanges;

    const body: Record<string, any> = { files: [file] };
    if (options?.modelVersion) body.model_version = options.modelVersion;
    if (options?.enableFormula !== undefined) body.enable_formula = options.enableFormula;
    if (options?.enableTable !== undefined) body.enable_table = options.enableTable;
    if (options?.language) body.language = options.language;

    const data = await fetchJson<MinerUBatchUrlResponse>(
      `${MINERU_URL}/api/v4/file-urls/batch`,
      {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          "Authorization": `Bearer ${apiKey}`,
        },
        body: JSON.stringify(body),
        signal,
      },
    );

    if (data.code !== 0) {
      return { batchId: "", uploadUrl: "", error: `MinerU error ${data.code}: ${data.msg}` };
    }

    const uploadUrl = data.data.file_urls?.[0];
    if (!uploadUrl) {
      return { batchId: "", uploadUrl: "", error: "MinerU did not return an upload URL" };
    }

    return { batchId: data.data.batch_id, uploadUrl };
  } catch (error) {
    if (isAbortError(error)) throw error;
    return { batchId: "", uploadUrl: "", error: formatServiceError("MinerU", error) };
  }
}

interface PollStatus {
  state: string;
  zipUrl?: string;
  errMsg?: string;
  progress?: { extracted_pages: number; total_pages: number };
}

async function fetchSingleStatus(taskId: string, apiKey: string, signal?: AbortSignal): Promise<PollStatus> {
  const data = await fetchJson<MinerUTaskResult>(
    `${MINERU_URL}/api/v4/extract/task/${taskId}`,
    { headers: { Authorization: `Bearer ${apiKey}` }, signal },
  );
  return {
    state: data.data.state,
    zipUrl: data.data.full_zip_url,
    errMsg: data.data.err_msg,
    progress: data.data.extract_progress,
  };
}

async function fetchBatchStatus(batchId: string, apiKey: string, signal?: AbortSignal): Promise<PollStatus> {
  const data = await fetchJson<MinerUBatchResult>(
    `${MINERU_URL}/api/v4/extract-results/batch/${batchId}`,
    { headers: { Authorization: `Bearer ${apiKey}` }, signal },
  );
  const item = data.data.extract_result?.[0];
  if (!item) return { state: "pending" };
  return {
    state: item.state,
    zipUrl: item.full_zip_url,
    errMsg: item.err_msg,
    progress: item.extract_progress,
  };
}

async function pollUntil(
  fetchStatus: () => Promise<PollStatus>,
  timeoutLabel: string,
  signal?: AbortSignal,
  onProgress?: (state: string, progress?: { extracted_pages: number; total_pages: number }) => void,
): Promise<{ zipUrl: string; error?: string }> {
  for (let attempt = 0; attempt < MAX_POLL_ATTEMPTS; attempt++) {
    if (signal?.aborted) throw new DOMException("Aborted", "AbortError");

    await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));

    try {
      const status = await fetchStatus();

      if (status.state === "done" && status.zipUrl) {
        return { zipUrl: status.zipUrl };
      }

      if (status.state === "failed") {
        return { zipUrl: "", error: status.errMsg || "MinerU parsing failed" };
      }

      if (status.state === "running" && status.progress) {
        onProgress?.(status.state, status.progress);
      } else {
        onProgress?.(status.state);
      }
    } catch (error) {
      if (isAbortError(error)) throw error;
      return { zipUrl: "", error: formatServiceError("MinerU", error) };
    }
  }

  return { zipUrl: "", error: `Timeout after ${MAX_POLL_ATTEMPTS * POLL_INTERVAL_MS / 1000}s. ${timeoutLabel}` };
}

async function downloadZip(zipUrl: string, signal?: AbortSignal): Promise<{ buffer?: Buffer; error?: string }> {
  const response = await fetch(zipUrl, { signal });
  if (!response.ok) {
    return { error: `Failed to download zip: HTTP ${response.status}` };
  }
  const arrayBuffer = await response.arrayBuffer();
  return { buffer: Buffer.from(arrayBuffer) };
}

// Walk the central directory (the authoritative ZIP entry index, avoids trusting
// local-header order) and return the raw bytes of the first entry whose name
// satisfies `match`. Handles stored (0) and deflate (8) entries; MinerU returns
// standard non-ZIP64 archives.
function extractZipEntry(buffer: Buffer, match: (name: string) => boolean): Buffer | null {
  let eocdOffset = -1;
  for (let i = buffer.length - 22; i >= Math.max(0, buffer.length - 65536); i--) {
    if (buffer[i] === 0x50 && buffer[i + 1] === 0x4b && buffer[i + 2] === 0x05 && buffer[i + 3] === 0x06) {
      eocdOffset = i;
      break;
    }
  }
  if (eocdOffset === -1) return null;

  const cdOffset = buffer.readUInt32LE(eocdOffset + 16);
  const cdEntries = buffer.readUInt16LE(eocdOffset + 10);

  let offset = cdOffset;
  for (let i = 0; i < cdEntries; i++) {
    if (buffer.readUInt32LE(offset) !== 0x02014b50) break;

    const compressionMethod = buffer.readUInt16LE(offset + 10);
    const compressedSize = buffer.readUInt32LE(offset + 20);
    const filenameLen = buffer.readUInt16LE(offset + 28);
    const extraLen = buffer.readUInt16LE(offset + 30);
    const commentLen = buffer.readUInt16LE(offset + 32);
    const localHeaderOffset = buffer.readUInt32LE(offset + 42);
    const filename = buffer.toString("utf8", offset + 46, offset + 46 + filenameLen);

    if (match(filename)) {
      const lhFilenameLen = buffer.readUInt16LE(localHeaderOffset + 26);
      const lhExtraLen = buffer.readUInt16LE(localHeaderOffset + 28);
      const dataOffset = localHeaderOffset + 30 + lhFilenameLen + lhExtraLen;
      const compressed = buffer.subarray(dataOffset, dataOffset + compressedSize);
      if (compressionMethod === 0) return compressed;
      if (compressionMethod === 8) return inflateRawSync(compressed);
      return null;
    }

    offset += 46 + filenameLen + extraLen + commentLen;
  }
  return null;
}

// Download the result zip once and extract either the flat full.md (default) or
// the page-grouped content_list_v2.json (for paginated rendering).
async function extractFromZip(
  zipUrl: string,
  signal: AbortSignal | undefined,
  paginate: boolean,
): Promise<{ markdown?: string; contentListV2?: unknown[]; error?: string }> {
  try {
    const { buffer, error } = await downloadZip(zipUrl, signal);
    if (error || !buffer) return { error: error || "Failed to download zip" };

    if (paginate) {
      const entry = extractZipEntry(buffer, (n) => n.endsWith("_content_list_v2.json"));
      if (!entry) return { error: "content_list_v2.json not found in ZIP archive" };
      try {
        return { contentListV2: JSON.parse(entry.toString("utf8")) };
      } catch {
        return { error: "Failed to parse content_list_v2.json" };
      }
    }

    const md = extractZipEntry(buffer, (n) => n === "full.md");
    if (!md) return { error: "full.md not found in ZIP archive" };
    return { markdown: md.toString("utf8") };
  } catch (error) {
    if (isAbortError(error)) throw error;
    return { error: formatServiceError("MinerU download", error) };
  }
}

export interface ParseResult {
  markdown: string;
  contentListV2?: unknown[];
  totalPages?: number;
  taskId?: string;
  error?: string;
}

export async function parsePdfWithMinerU(
  url: string,
  apiKey: string,
  signal?: AbortSignal,
  onProgress?: (phase: string, detail?: string) => void,
  options?: ExtractOptions,
  paginate = false,
): Promise<ParseResult> {
  // Step 1: Submit task
  onProgress?.("submitting", url.length > 80 ? url.slice(0, 77) + "..." : url);
  const { taskId, error: submitError } = await submitTask(url, apiKey, signal, options);
  if (submitError || !taskId) {
    return { markdown: "", error: submitError || "Failed to get task ID" };
  }

  // Step 2: Poll until done
  let totalPages: number | undefined;
  onProgress?.("parsing", `task ${taskId.slice(0, 8)}...`);
  const { zipUrl, error: pollError } = await pollUntil(
    () => fetchSingleStatus(taskId, apiKey, signal),
    `Task ID: ${taskId}`,
    signal,
    (state, progress) => {
      if (progress) totalPages = progress.total_pages;
      if (state === "running" && progress) {
        onProgress?.("parsing", `${progress.extracted_pages}/${progress.total_pages} pages`);
      } else {
        onProgress?.("parsing", state);
      }
    },
  );

  if (pollError || !zipUrl) {
    // Include taskId so user can check later
    return { markdown: "", taskId, totalPages, error: pollError || "Failed to get result" };
  }

  // Step 3: Download zip and extract result (full.md, or content_list_v2 when paginating)
  onProgress?.("downloading", "fetching results...");
  const ex = await extractFromZip(zipUrl, signal, paginate);
  if (ex.error) {
    return { markdown: "", taskId, totalPages, error: ex.error };
  }

  return { markdown: ex.markdown ?? "", contentListV2: ex.contentListV2, totalPages, taskId };
}

export async function parseLocalFileWithMinerU(
  filePath: string,
  apiKey: string,
  signal?: AbortSignal,
  onProgress?: (phase: string, detail?: string) => void,
  options?: ExtractOptions,
  paginate = false,
): Promise<ParseResult> {
  const fileName = basename(filePath);

  // Step 1: Request a signed upload URL (batch endpoint, single file)
  onProgress?.("submitting", fileName);
  const { batchId, uploadUrl, error: applyError } = await requestUploadUrl(fileName, apiKey, signal, options);
  if (applyError || !batchId || !uploadUrl) {
    return { markdown: "", error: applyError || "Failed to get upload URL" };
  }

  // Step 2: PUT the file to the signed OSS URL — no auth header, no content-type (signature is in the URL)
  onProgress?.("uploading", fileName);
  try {
    const fileBuffer = await readFile(filePath);
    const putResponse = await fetch(uploadUrl, { method: "PUT", body: fileBuffer, signal });
    if (!putResponse.ok) {
      return { markdown: "", taskId: batchId, error: `Upload failed: HTTP ${putResponse.status}` };
    }
  } catch (error) {
    if (isAbortError(error)) throw error;
    return { markdown: "", taskId: batchId, error: formatServiceError("MinerU upload", error) };
  }

  // Step 3: Poll batch results until done (upload auto-submits the task)
  let totalPages: number | undefined;
  onProgress?.("parsing", `batch ${batchId.slice(0, 8)}...`);
  const { zipUrl, error: pollError } = await pollUntil(
    () => fetchBatchStatus(batchId, apiKey, signal),
    `Batch ID: ${batchId}`,
    signal,
    (state, progress) => {
      if (progress) totalPages = progress.total_pages;
      if (state === "running" && progress) {
        onProgress?.("parsing", `${progress.extracted_pages}/${progress.total_pages} pages`);
      } else {
        onProgress?.("parsing", state);
      }
    },
  );

  if (pollError || !zipUrl) {
    return { markdown: "", taskId: batchId, totalPages, error: pollError || "Failed to get result" };
  }

  // Step 4: Download zip and extract result (shared with the URL path)
  onProgress?.("downloading", "fetching results...");
  const ex = await extractFromZip(zipUrl, signal, paginate);
  if (ex.error) {
    return { markdown: "", taskId: batchId, totalPages, error: ex.error };
  }

  return { markdown: ex.markdown ?? "", contentListV2: ex.contentListV2, totalPages, taskId: batchId };
}