#!/usr/bin/env python3
"""Upload one local document to a MinerU-compatible v4 API and download its results."""

from __future__ import annotations

import argparse
import http.client
import json
import os
import re
import shutil
import stat
import sys
import tempfile
import time
import uuid
import zipfile
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import quote, urlsplit
from urllib.request import Request, urlopen


DEFAULT_SUBMIT_ENDPOINT = "https://mineru.net/api/v4/file-urls/batch"
DEFAULT_RESULT_ENDPOINT_TEMPLATE = (
    "https://mineru.net/api/v4/extract-results/batch/{batch_id}"
)
TERMINAL_STATES = {"done", "failed"}
SEARCHABLE_SUFFIXES = {".html", ".json", ".md", ".tex", ".txt"}


class MineruError(RuntimeError):
    """Raised for API, transport, archive, or verification failures."""


def _redact(value: str, token: str) -> str:
    return value.replace(token, "<redacted>") if token else value


def _api_request(
    method: str,
    url: str,
    token: str,
    timeout: float,
    payload: dict[str, Any] | None = None,
) -> dict[str, Any]:
    body = None if payload is None else json.dumps(payload).encode("utf-8")
    headers = {"Accept": "application/json"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    if body is not None:
        headers["Content-Type"] = "application/json"

    request = Request(url, data=body, headers=headers, method=method)
    try:
        with urlopen(request, timeout=timeout) as response:
            raw = response.read()
    except HTTPError as error:
        detail = error.read(8192).decode("utf-8", errors="replace")
        raise MineruError(
            _redact(f"MinerU API returned HTTP {error.code}: {detail}", token)
        ) from error
    except (URLError, OSError) as error:
        raise MineruError(f"MinerU API request failed: {error.reason}") from error

    try:
        result = json.loads(raw)
    except json.JSONDecodeError as error:
        raise MineruError("MinerU API returned invalid JSON") from error
    if not isinstance(result, dict):
        raise MineruError("MinerU API returned a non-object response")
    if result.get("code") != 0:
        message = _redact(str(result.get("msg", "unknown API error")), token)
        trace_id = result.get("trace_id")
        suffix = f" (trace_id={trace_id})" if trace_id else ""
        raise MineruError(f"MinerU API rejected the request: {message}{suffix}")
    data = result.get("data")
    if not isinstance(data, dict):
        raise MineruError("MinerU API response is missing its data object")
    return data


def _put_file(upload_url: str, file_path: Path, timeout: float) -> None:
    parsed = urlsplit(upload_url)
    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
        raise MineruError("MinerU returned an invalid upload URL")

    connection_class = (
        http.client.HTTPSConnection if parsed.scheme == "https" else http.client.HTTPConnection
    )
    connection = connection_class(parsed.hostname, parsed.port, timeout=timeout)
    target = parsed.path or "/"
    if parsed.query:
        target = f"{target}?{parsed.query}"

    try:
        connection.putrequest("PUT", target)
        connection.putheader("Content-Length", str(file_path.stat().st_size))
        connection.endheaders()
        with file_path.open("rb") as source:
            while chunk := source.read(1024 * 1024):
                connection.send(chunk)
        response = connection.getresponse()
        detail = response.read(8192).decode("utf-8", errors="replace")
        if not 200 <= response.status < 300:
            raise MineruError(
                f"MinerU file upload returned HTTP {response.status}: {detail}"
            )
    except (OSError, http.client.HTTPException) as error:
        raise MineruError(f"MinerU file upload failed: {error}") from error
    finally:
        connection.close()


def _download(url: str, destination: Path, timeout: float) -> None:
    request = Request(url, headers={"Accept": "application/zip"})
    try:
        with urlopen(request, timeout=timeout) as response, destination.open("wb") as output:
            shutil.copyfileobj(response, output, length=1024 * 1024)
    except (HTTPError, URLError, OSError) as error:
        raise MineruError(f"Failed to download MinerU result archive: {error}") from error


def _safe_extract(archive_path: Path, output_dir: Path) -> None:
    output_root = output_dir.resolve()

    try:
        with zipfile.ZipFile(archive_path) as archive:
            validated_members: list[tuple[zipfile.ZipInfo, Path]] = []
            for member in archive.infolist():
                mode = member.external_attr >> 16
                if stat.S_ISLNK(mode):
                    raise MineruError(f"Result archive contains a symbolic link: {member.filename}")

                target = (output_dir / member.filename).resolve()
                try:
                    target.relative_to(output_root)
                except ValueError as error:
                    raise MineruError(
                        f"Result archive contains an unsafe path: {member.filename}"
                    ) from error
                validated_members.append((member, target))

            output_dir.mkdir(parents=True, exist_ok=False)
            for member, target in validated_members:
                if member.is_dir():
                    target.mkdir(parents=True, exist_ok=True)
                    continue
                target.parent.mkdir(parents=True, exist_ok=True)
                with archive.open(member) as source, target.open("wb") as destination:
                    shutil.copyfileobj(source, destination, length=1024 * 1024)
    except zipfile.BadZipFile as error:
        raise MineruError("MinerU result is not a valid ZIP archive") from error


def _make_data_id(file_path: Path) -> str:
    stem = re.sub(r"[^A-Za-z0-9_.-]+", "-", file_path.stem).strip("-.") or "document"
    return f"{stem[:100]}-{uuid.uuid4().hex[:12]}"


def _make_payload(args: argparse.Namespace, data_id: str) -> dict[str, Any]:
    file_entry: dict[str, Any] = {
        "name": args.input.name,
        "data_id": data_id,
        "is_ocr": args.ocr,
    }
    if args.pages:
        file_entry["page_ranges"] = args.pages

    payload: dict[str, Any] = {
        "files": [file_entry],
        "model_version": args.model,
        "language": args.language,
        "enable_formula": args.formula,
        "enable_table": args.table,
    }
    if args.extra_format:
        payload["extra_formats"] = args.extra_format
    return payload


def _validate_endpoint(endpoint: str, name: str) -> str:
    parsed = urlsplit(endpoint)
    if not parsed.hostname or parsed.username or parsed.password:
        raise MineruError(f"Invalid MinerU {name}")
    is_local = parsed.hostname in {"127.0.0.1", "::1", "localhost"}
    if parsed.scheme != "https" and not (parsed.scheme == "http" and is_local):
        raise MineruError(f"MinerU {name} must use HTTPS unless it is localhost")
    if parsed.fragment:
        raise MineruError(f"MinerU {name} must not contain a URL fragment")
    return endpoint


def _find_extract_result(data: dict[str, Any], data_id: str) -> dict[str, Any]:
    results = data.get("extract_result")
    if not isinstance(results, list) or not results:
        raise MineruError("MinerU batch result does not contain extract_result")
    matches = [entry for entry in results if isinstance(entry, dict) and entry.get("data_id") == data_id]
    if len(matches) == 1:
        return matches[0]
    if len(results) == 1 and isinstance(results[0], dict):
        return results[0]
    raise MineruError("MinerU batch result could not be matched to the uploaded file")


def _poll_result(
    result_endpoint_template: str,
    token: str,
    batch_id: str,
    data_id: str,
    poll_interval: float,
    overall_timeout: float,
    request_timeout: float,
) -> str:
    deadline = time.monotonic() + overall_timeout
    previous_progress: tuple[object, object, object] | None = None
    result_url = result_endpoint_template.replace(
        "{batch_id}", quote(batch_id, safe="")
    )

    while time.monotonic() < deadline:
        data = _api_request("GET", result_url, token, request_timeout)
        result = _find_extract_result(data, data_id)
        state = result.get("state")
        progress = result.get("extract_progress")
        extracted = progress.get("extracted_pages") if isinstance(progress, dict) else None
        total = progress.get("total_pages") if isinstance(progress, dict) else None
        progress_key = (state, extracted, total)
        if progress_key != previous_progress:
            detail = f" {extracted}/{total} pages" if extracted is not None and total is not None else ""
            print(f"MinerU state: {state}{detail}")
            previous_progress = progress_key

        if state == "done":
            zip_url = result.get("full_zip_url")
            if not isinstance(zip_url, str) or not zip_url:
                raise MineruError("Completed MinerU task is missing full_zip_url")
            return zip_url
        if state == "failed":
            raise MineruError(f"MinerU parsing failed: {result.get('err_msg', 'unknown error')}")
        if not isinstance(state, str) or state in TERMINAL_STATES:
            raise MineruError(f"MinerU returned an invalid task state: {state!r}")
        time.sleep(min(poll_interval, max(0.0, deadline - time.monotonic())))

    raise MineruError(f"MinerU task did not finish within {overall_timeout:g} seconds")


def _verify_expectations(output_dir: Path, expectations: list[str]) -> None:
    if not expectations:
        return

    searchable_files = sorted(
        path
        for path in output_dir.rglob("*")
        if path.is_file() and path.suffix.lower() in SEARCHABLE_SUFFIXES
    )
    missing: list[str] = []
    for expectation in expectations:
        found_in: list[str] = []
        for path in searchable_files:
            try:
                content = path.read_text(encoding="utf-8", errors="replace")
            except OSError as error:
                raise MineruError(f"Failed to inspect {path}: {error}") from error
            if expectation in content:
                found_in.append(str(path.relative_to(output_dir)))
        if found_in:
            print(f"Verified {expectation!r} in: {', '.join(found_in)}")
        else:
            print(f"Missing expected text: {expectation!r}", file=sys.stderr)
            missing.append(expectation)
    if missing:
        raise MineruError(f"Result verification failed for {len(missing)} expected value(s)")


def _default_output(input_path: Path) -> Path:
    return Path("output") / "mineru" / input_path.stem


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Upload one local document to MinerU's precise parsing API."
    )
    parser.add_argument("input", type=Path, help="local PDF or other supported document")
    parser.add_argument("-o", "--output", type=Path, help="new directory for extracted results")
    parser.add_argument("--model", choices=("pipeline", "vlm"), default="pipeline")
    parser.add_argument("--language", default="ch")
    parser.add_argument("--pages", help='page range such as "2,4-6"')
    parser.add_argument("--ocr", action=argparse.BooleanOptionalAction, default=False)
    parser.add_argument("--formula", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--table", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument(
        "--extra-format",
        action="append",
        choices=("docx", "html", "latex"),
        default=[],
    )
    parser.add_argument(
        "--expect",
        action="append",
        default=[],
        help="exact text that must occur in an extracted text/JSON file; repeatable",
    )
    parser.add_argument("--poll-interval", type=float, default=5.0)
    parser.add_argument("--timeout", type=float, default=1800.0, help="overall polling timeout")
    parser.add_argument("--request-timeout", type=float, default=60.0)
    parser.add_argument("--keep-zip", action="store_true")
    parser.add_argument(
        "--allow-anonymous",
        action="store_true",
        help="allow an API request without a bearer token, intended for a local service",
    )
    parser.add_argument("--dry-run", action="store_true", help="print the request payload only")
    parser.add_argument(
        "--submit-endpoint",
        default=DEFAULT_SUBMIT_ENDPOINT,
        help="endpoint that creates a local-file upload batch",
    )
    parser.add_argument(
        "--result-endpoint-template",
        default=DEFAULT_RESULT_ENDPOINT_TEMPLATE,
        help="polling endpoint containing the {batch_id} placeholder",
    )
    return parser


def _validate_args(args: argparse.Namespace) -> None:
    args.input = args.input.expanduser().resolve()
    if not args.input.is_file():
        raise MineruError(f"Input file does not exist: {args.input}")
    if args.poll_interval <= 0 or args.timeout <= 0 or args.request_timeout <= 0:
        raise MineruError("Timeout and poll interval values must be positive")
    args.output = (args.output or _default_output(args.input)).expanduser().resolve()
    if args.output.exists():
        raise MineruError(f"Output path already exists: {args.output}")
    args.submit_endpoint = _validate_endpoint(args.submit_endpoint, "submit endpoint")
    if args.result_endpoint_template.count("{batch_id}") != 1:
        raise MineruError("MinerU result endpoint template must contain {batch_id} exactly once")
    args.result_endpoint_template = _validate_endpoint(
        args.result_endpoint_template,
        "result endpoint template",
    )


def run(args: argparse.Namespace) -> Path:
    _validate_args(args)
    data_id = _make_data_id(args.input)
    payload = _make_payload(args, data_id)
    if args.dry_run:
        print(json.dumps(payload, ensure_ascii=False, indent=2))
        return args.output

    token = os.environ.get("MINERU_TOKEN", "").strip()
    if not token and not args.allow_anonymous:
        raise MineruError("MINERU_TOKEN is not set")

    submitted = _api_request(
        "POST",
        args.submit_endpoint,
        token,
        args.request_timeout,
        payload,
    )
    batch_id = submitted.get("batch_id")
    upload_urls = submitted.get("file_urls")
    if not isinstance(batch_id, str) or not batch_id:
        raise MineruError("MinerU upload response is missing batch_id")
    if not isinstance(upload_urls, list) or len(upload_urls) != 1 or not isinstance(upload_urls[0], str):
        raise MineruError("MinerU upload response must contain exactly one file URL")

    print(f"MinerU batch created: {batch_id}")
    _put_file(upload_urls[0], args.input, args.request_timeout)
    print(f"Uploaded: {args.input.name}")

    zip_url = _poll_result(
        args.result_endpoint_template,
        token,
        batch_id,
        data_id,
        args.poll_interval,
        args.timeout,
        args.request_timeout,
    )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    temporary_zip: Path | None = None
    try:
        with tempfile.NamedTemporaryFile(
            prefix=f"{args.input.stem}-",
            suffix=".zip",
            dir=args.output.parent,
            delete=False,
        ) as temporary:
            temporary_zip = Path(temporary.name)
        _download(zip_url, temporary_zip, args.request_timeout)
        _safe_extract(temporary_zip, args.output)
        _verify_expectations(args.output, args.expect)
        if args.keep_zip:
            shutil.move(str(temporary_zip), args.output / "mineru-result.zip")
            temporary_zip = None
    finally:
        if temporary_zip is not None:
            temporary_zip.unlink(missing_ok=True)

    print(f"MinerU results: {args.output}")
    return args.output


def main(argv: list[str] | None = None) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)
    try:
        run(args)
    except MineruError as error:
        print(f"error: {error}", file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        print("error: interrupted", file=sys.stderr)
        return 130
    return 0


if __name__ == "__main__":
    raise SystemExit(main())