#!/usr/bin/env python3
"""OCR script: convert images/PDFs to markdown via PaddleOCR layout-parsing API."""

import argparse
import base64
import json
import sys
import urllib.request
from pathlib import Path

API_BASE = "http://192.168.25.212:8080"

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".webp"}
PDF_EXTS = {".pdf"}


def encode_file(file_path: Path) -> tuple[str, int]:
    """Read a file and return (base64_string, file_type). file_type: 0=PDF, 1=image."""
    ext = file_path.suffix.lower()
    if ext in PDF_EXTS:
        file_type = 0
    elif ext in IMAGE_EXTS:
        file_type = 1
    else:
        sys.exit(f"Unsupported file type: {ext}")

    with open(file_path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode("utf-8")
    return b64, file_type


def call_api(
    file_value: str,
    file_type: int | None = None,
    use_layout_detection: bool = True,
    prettify_markdown: bool = True,
    restructure_pages: bool = True,
    relevel_titles: bool = True,
    merge_tables: bool = True,
) -> dict:
    """Call the layout-parsing API and return the JSON response."""
    payload: dict = {
        "file": file_value,
        "useLayoutDetection": use_layout_detection,
        "prettifyMarkdown": prettify_markdown,
        "restructurePages": restructure_pages,
        "relevelTitles": relevel_titles,
        "mergeTables": merge_tables,
        "visualize": False,
    }
    if file_type is not None:
        payload["fileType"] = file_type

    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        f"{API_BASE}/layout-parsing",
        data=data,
        headers={"Content-Type": "application/json"},
    )

    try:
        with urllib.request.urlopen(req, timeout=300) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.URLError as e:
        sys.exit(f"API request failed: {e}")


def extract_markdown(response: dict) -> str:
    """Extract and concatenate markdown text from the API response."""
    if response.get("errorCode", -1) != 0:
        sys.exit(f"API error: {response.get('errorMsg', 'Unknown error')}")

    results = response.get("result", {}).get("layoutParsingResults", [])
    if not results:
        sys.exit("No results returned from API.")

    pages = []
    for i, page in enumerate(results):
        md = page.get("markdown", {}).get("text", "")
        if md:
            pages.append(md)

    return "\n\n---\n\n".join(pages)


def determine_output_path(input_path: str, output: str | None) -> Path:
    """Determine the output .md file path."""
    if output:
        return Path(output)
    p = Path(input_path)
    return p.with_suffix(".md")


def main():
    parser = argparse.ArgumentParser(description="OCR: image/PDF to Markdown")
    parser.add_argument("input", help="Local file path or URL")
    parser.add_argument("-o", "--output", help="Output markdown file path (default: same name with .md extension)")
    parser.add_argument("--no-layout", action="store_true", help="Disable layout detection")
    parser.add_argument("--no-restructure", action="store_true", help="Disable page restructuring for multi-page PDFs")
    args = parser.parse_args()

    input_value = args.input
    is_url = input_value.startswith("http://") or input_value.startswith("https://")

    if is_url:
        # Pass URL directly to the API
        file_value = input_value
        file_type = None
        if args.output:
            out_path = Path(args.output)
        else:
            # Derive name from URL
            from urllib.parse import urlparse
            url_path = urlparse(input_value).path
            name = Path(url_path).stem or "output"
            out_path = Path(f"{name}.md")
    else:
        # Local file: base64 encode
        file_path = Path(input_value).expanduser().resolve()
        if not file_path.exists():
            sys.exit(f"File not found: {file_path}")
        file_value, file_type = encode_file(file_path)
        out_path = determine_output_path(str(file_path), args.output)

    response = call_api(
        file_value=file_value,
        file_type=file_type,
        use_layout_detection=not args.no_layout,
        restructure_pages=not args.no_restructure,
    )

    markdown_text = extract_markdown(response)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(markdown_text, encoding="utf-8")
    print(out_path.resolve())


if __name__ == "__main__":
    main()
