#!/usr/bin/env python3
"""Download an official Chinese national standard and parse it with MinerU."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
import uuid
from pathlib import Path
from typing import NamedTuple
from urllib.parse import quote, urlsplit


STANDARD_PATTERN = re.compile(
    r"^\s*(GB(?:\s*/\s*[TZ])?)\s*([0-9]+(?:\.[0-9]+)*)\s*-\s*([0-9]{4})\s*$",
    re.IGNORECASE,
)
PAGES_PATTERN = re.compile(r"^-?\d+(?:--?\d+)?(?:,-?\d+(?:--?\d+)?)*$")
DEFAULT_OUTPUT_ROOT = Path("china-standard-output")
DEFAULT_CONFIG_PATH = Path(__file__).resolve().parents[1] / "config" / "mineru.json"
MANIFEST_SCHEMA_VERSION = 1
PARSE_PROFILE = {
    "profile_version": 1,
    "model_version": "vlm",
    "is_ocr": True,
    "enable_formula": False,
    "enable_table": True,
    "language": "ch",
    "extra_formats": ["docx", "html", "latex"],
}
RESERVED_OUTPUT_NAMES = {"source.pdf", "cache-manifest.json", ".parse.lock"}


class StandardParseError(RuntimeError):
    """Raised for validation, browser automation, or MinerU failures."""


class MineruServiceConfig(NamedTuple):
    """Validated MinerU service connection settings."""

    path: Path
    mode: str
    submit_endpoint: str
    result_endpoint_template: str
    token: str
    token_source: str
    request_timeout: float
    parse_timeout: float
    poll_interval: float


def normalize_standard_number(value: str) -> str:
    match = STANDARD_PATTERN.fullmatch(value)
    if match is None:
        raise StandardParseError(
            "Invalid standard number. Expected GB, GB/T, or GB/Z followed by a number and four-digit year."
        )
    prefix = re.sub(r"\s+", "", match.group(1).upper())
    return f"{prefix} {match.group(2)}-{match.group(3)}"


def standard_slug(standard_number: str) -> str:
    return standard_number.replace("/", "").replace(" ", "")


def validate_pages(value: str | None) -> str | None:
    if value is None:
        return None
    compact = re.sub(r"\s+", "", value)
    if not compact or PAGES_PATTERN.fullmatch(compact) is None:
        raise StandardParseError('Invalid page range. Use a value such as "2,4-6".')
    return compact


def find_playwright_cli() -> str:
    names = ("playwright-cli.cmd", "playwright-cli") if os.name == "nt" else ("playwright-cli",)
    package_bin = Path(__file__).resolve().parents[1] / "node_modules" / ".bin"
    for name in names:
        candidate = package_bin / name
        if candidate.is_file():
            return str(candidate)
    for name in names:
        executable = shutil.which(name)
        if executable:
            return executable
    raise StandardParseError(
        "playwright-cli is unavailable. Install @playwright/cli or make playwright-cli available on PATH."
    )


def run_command(
    command: list[str],
    timeout: float,
    token: str = "",
    environment: dict[str, str] | None = None,
) -> subprocess.CompletedProcess[str]:
    try:
        result = subprocess.run(
            command,
            check=False,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=timeout,
            env=environment,
        )
    except (OSError, subprocess.TimeoutExpired) as error:
        raise StandardParseError(f"Command failed: {error}") from error

    if result.returncode != 0:
        detail = "\n".join(part.strip() for part in (result.stdout, result.stderr) if part.strip())
        if token:
            detail = detail.replace(token, "<redacted>")
        raise StandardParseError(f"Command returned {result.returncode}: {detail[-8000:]}")
    return result


def positive_number(data: dict[str, object], key: str, default: float) -> float:
    value = data.get(key, default)
    if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0:
        raise StandardParseError(f"{key} must be a positive number")
    return float(value)


def validate_service_endpoint(
    endpoint: object,
    field: str,
    mode: str,
    require_batch_id: bool,
) -> str:
    if not isinstance(endpoint, str) or not endpoint.strip():
        raise StandardParseError(f"{field} must be a non-empty URL string")
    endpoint = endpoint.strip()
    if require_batch_id and endpoint.count("{batch_id}") != 1:
        raise StandardParseError(f"{field} must contain {{batch_id}} exactly once")
    if not require_batch_id and "{batch_id}" in endpoint:
        raise StandardParseError(f"{field} must not contain {{batch_id}}")
    try:
        parsed = urlsplit(endpoint)
        parsed_port = parsed.port
    except ValueError as error:
        raise StandardParseError(f"{field} is invalid") from error
    if (
        parsed.scheme not in {"http", "https"}
        or parsed.hostname is None
        or parsed.username is not None
        or parsed.password is not None
        or parsed_port is not None and not 1 <= parsed_port <= 65535
        or parsed.fragment
    ):
        raise StandardParseError(f"{field} is invalid")
    loopback_hosts = {"127.0.0.1", "localhost", "::1"}
    if mode == "cloud" and parsed.scheme != "https":
        raise StandardParseError(f"{field} must use HTTPS in cloud mode")
    if mode == "local" and parsed.scheme == "http" and parsed.hostname.lower() not in loopback_hosts:
        raise StandardParseError(f"{field} may use HTTP only for a loopback host")
    return endpoint


def load_service_config(path: Path) -> MineruServiceConfig:
    resolved = path.expanduser().resolve()
    try:
        data = json.loads(resolved.read_text(encoding="utf-8"))
    except FileNotFoundError as error:
        raise StandardParseError(
            f"MinerU configuration is missing: {resolved}. Copy mineru.example.json to mineru.json."
        ) from error
    except (OSError, json.JSONDecodeError) as error:
        raise StandardParseError(f"Cannot read MinerU configuration: {error}") from error
    if not isinstance(data, dict):
        raise StandardParseError("MinerU configuration must be a JSON object")

    mode = data.get("mode")
    if mode not in {"local", "cloud"}:
        raise StandardParseError('mode must be either "local" or "cloud"')

    default_submit_endpoint = (
        "http://127.0.0.1:8000/api/v4/file-urls/batch"
        if mode == "local"
        else "https://mineru.net/api/v4/file-urls/batch"
    )
    default_result_endpoint_template = (
        "http://127.0.0.1:8000/api/v4/extract-results/batch/{batch_id}"
        if mode == "local"
        else "https://mineru.net/api/v4/extract-results/batch/{batch_id}"
    )
    submit_endpoint = validate_service_endpoint(
        data.get("submit_endpoint", default_submit_endpoint),
        "submit_endpoint",
        mode,
        require_batch_id=False,
    )
    result_endpoint_template = validate_service_endpoint(
        data.get("result_endpoint_template", default_result_endpoint_template),
        "result_endpoint_template",
        mode,
        require_batch_id=True,
    )
    token_setting = data.get("token", "")

    if not isinstance(token_setting, str):
        raise StandardParseError("token must be a string")
    token_setting = token_setting.strip()
    if token_setting.startswith("env:"):
        variable = token_setting.removeprefix("env:").strip()
        if not variable or re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", variable) is None:
            raise StandardParseError("token environment reference must look like env:MINERU_TOKEN")
        token = os.environ.get(variable, "").strip()
        token_source = f"env:{variable}"
    else:
        token = token_setting
        token_source = "config" if token else "none"
    request_timeout = positive_number(data, "request_timeout_seconds", 60.0)
    parse_timeout = positive_number(data, "parse_timeout_seconds", 1800.0)
    poll_interval = positive_number(data, "poll_interval_seconds", 5.0)
    return MineruServiceConfig(
        path=resolved,
        mode=mode,
        submit_endpoint=submit_endpoint,
        result_endpoint_template=result_endpoint_template,
        token=token,
        token_source=token_source,
        request_timeout=request_timeout,
        parse_timeout=parse_timeout,
        poll_interval=poll_interval,
    )


def make_browser_program(standard_number: str, pdf_path: Path) -> str:
    search_url = (
        "https://openstd.samr.gov.cn/bzgk/std/std_list?"
        "p.p1=0&p.p90=circulation_date&p.p91=desc&p.p2="
        f"{quote(standard_number)}"
    )
    standard_json = json.dumps(standard_number, ensure_ascii=False)
    url_json = json.dumps(search_url, ensure_ascii=False)
    path_json = json.dumps(str(pdf_path.resolve()), ensure_ascii=False)
    return f"""async page => {{
  const standardNumber = {standard_json};
  const searchUrl = {url_json};
  const outputPath = {path_json};
  await page.goto(searchUrl, {{ waitUntil: 'domcontentloaded', timeout: 60000 }});

  const exactLink = page.getByRole('link', {{ name: standardNumber, exact: true }});
  const resultRows = page.getByRole('row').filter({{ has: exactLink }});
  const resultCount = await resultRows.count();
  if (resultCount !== 1) {{
    throw new Error(`Expected one exact result for ${{standardNumber}}, found ${{resultCount}}`);
  }}

  const detailPagePromise = page.context().waitForEvent('page', {{ timeout: 30000 }});
  await resultRows.first().getByRole('button', {{ name: '查看详细', exact: true }}).click();
  const detailPage = await detailPagePromise;
  await detailPage.waitForLoadState('domcontentloaded', {{ timeout: 60000 }});

  const heading = detailPage.getByRole('heading', {{ name: `标准号：${{standardNumber}}`, exact: true }});
  if (await heading.count() !== 1) {{
    throw new Error(`Detail page does not match ${{standardNumber}}`);
  }}

  const downloadButton = detailPage.getByRole('button', {{ name: '下载标准', exact: true }});
  if (await downloadButton.count() !== 1) {{
    throw new Error(`Official PDF is not publicly downloadable for ${{standardNumber}}`);
  }}

  const downloadPromise = detailPage.context().waitForEvent('page', {{ timeout: 30000 }}).then(
    downloadPage => downloadPage.waitForEvent('download', {{ timeout: 60000 }})
  );
  await downloadButton.click();
  const download = await downloadPromise;
  await download.saveAs(outputPath);
  return {{
    standard_number: standardNumber,
    detail_url: detailPage.url(),
    suggested_filename: download.suggestedFilename(),
    pdf: outputPath
  }};
}}"""


def download_standard(
    standard_number: str,
    pdf_path: Path,
    command_timeout: float,
) -> dict[str, object]:
    executable = find_playwright_cli()
    session = f"gb-{os.getpid()}-{uuid.uuid4().hex[:8]}"
    program_path: Path | None = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="w",
            encoding="utf-8",
            suffix=".js",
            prefix="download-china-standard-",
            delete=False,
        ) as program:
            program.write(make_browser_program(standard_number, pdf_path))
            program_path = Path(program.name)

        run_command([executable, f"-s={session}", "open", "about:blank"], command_timeout)
        result = run_command(
            [
                executable,
                f"-s={session}",
                "--raw",
                "run-code",
                f"--filename={program_path}",
            ],
            command_timeout,
        )
        try:
            browser_result = json.loads(result.stdout)
        except json.JSONDecodeError as error:
            raise StandardParseError("playwright-cli returned invalid download metadata") from error
        if not isinstance(browser_result, dict):
            raise StandardParseError("playwright-cli returned non-object download metadata")
        return browser_result
    finally:
        if program_path is not None:
            program_path.unlink(missing_ok=True)
        try:
            subprocess.run(
                [executable, f"-s={session}", "close"],
                check=False,
                capture_output=True,
                timeout=15,
            )
        except (OSError, subprocess.TimeoutExpired):
            pass


def validate_pdf(path: Path) -> None:
    try:
        size = path.stat().st_size
        with path.open("rb") as source:
            header = source.read(5)
    except OSError as error:
        raise StandardParseError(f"Cannot read downloaded PDF: {error}") from error
    if size < 1024 or header != b"%PDF-":
        raise StandardParseError("Official-site download is not a valid PDF")


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    try:
        with path.open("rb") as source:
            while chunk := source.read(1024 * 1024):
                digest.update(chunk)
    except OSError as error:
        raise StandardParseError(f"Cannot hash cached file: {error}") from error
    return digest.hexdigest()


def parse_profile(pages: str | None) -> dict[str, object]:
    return {**PARSE_PROFILE, "pages": pages}


def parse_profile_sha256(pages: str | None) -> str:
    encoded = json.dumps(
        parse_profile(pages),
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
    ).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def load_cache_manifest(path: Path) -> dict[str, object] | None:
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (FileNotFoundError, OSError, json.JSONDecodeError):
        return None
    if not isinstance(data, dict):
        return None
    return data


def write_cache_manifest(path: Path, manifest: dict[str, object]) -> None:
    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
    try:
        temporary.write_text(
            json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
        os.replace(temporary, path)
    except OSError as error:
        raise StandardParseError(f"Cannot write cache manifest: {error}") from error
    finally:
        temporary.unlink(missing_ok=True)


def cached_source_sha256(
    manifest: dict[str, object] | None,
    standard_number: str,
    pdf_path: Path,
) -> str | None:
    if (
        manifest is None
        or manifest.get("schema_version") != MANIFEST_SCHEMA_VERSION
        or manifest.get("standard_number") != standard_number
    ):
        return None
    source = manifest.get("source")
    if not isinstance(source, dict) or source.get("path") != pdf_path.name:
        return None
    expected_sha256 = source.get("sha256")
    if not isinstance(expected_sha256, str) or not expected_sha256:
        return None
    try:
        validate_pdf(pdf_path)
        actual_sha256 = file_sha256(pdf_path)
    except StandardParseError:
        return None
    return actual_sha256 if actual_sha256 == expected_sha256 else None


def cached_parse_outputs(
    manifest: dict[str, object],
    run_directory: Path,
    source_sha256: str,
    profile_sha256: str,
) -> list[str] | None:
    parsed = manifest.get("parse")
    if not isinstance(parsed, dict):
        return None
    if (
        parsed.get("source_sha256") != source_sha256
        or parsed.get("profile_sha256") != profile_sha256
    ):
        return None
    outputs = parsed.get("outputs")
    if not isinstance(outputs, list) or "full.md" not in outputs:
        return None
    resolved_root = run_directory.resolve()
    validated: list[str] = []
    for value in outputs:
        if not isinstance(value, str) or not value:
            return None
        candidate = (run_directory / value).resolve()
        try:
            candidate.relative_to(resolved_root)
        except ValueError:
            return None
        if not candidate.is_file() or candidate.stat().st_size == 0:
            return None
        validated.append(value)
    return validated


def find_mineru_client() -> Path:
    client = Path(__file__).with_name("mineru_cloud.py")
    if not client.is_file():
        raise StandardParseError(f"MinerU client not found: {client}")
    return client


def parse_with_mineru(
    pdf_path: Path,
    mineru_output: Path,
    pages: str | None,
    service: MineruServiceConfig,
) -> None:
    command = [
        sys.executable,
        str(find_mineru_client()),
        str(pdf_path),
        "--output",
        str(mineru_output),
        "--model",
        "vlm",
        "--ocr",
        "--no-formula",
        "--table",
        "--language",
        "ch",
        "--extra-format",
        "docx",
        "--extra-format",
        "html",
        "--extra-format",
        "latex",
        "--keep-zip",
        "--timeout",
        str(service.parse_timeout),
        "--request-timeout",
        str(service.request_timeout),
        "--poll-interval",
        str(service.poll_interval),
        "--submit-endpoint",
        service.submit_endpoint,
        "--result-endpoint-template",
        service.result_endpoint_template,
    ]
    if service.mode == "local" and not service.token:
        command.append("--allow-anonymous")
    if pages is not None:
        command.extend(("--pages", pages))
    environment = os.environ.copy()
    if service.token:
        environment["MINERU_TOKEN"] = service.token
    else:
        environment.pop("MINERU_TOKEN", None)
    result = run_command(
        command,
        service.parse_timeout + 120,
        service.token,
        environment,
    )
    if result.stdout:
        print(result.stdout.rstrip(), file=sys.stderr)
    if result.stderr:
        print(result.stderr.rstrip(), file=sys.stderr)


def staged_outputs(staging_directory: Path) -> list[str]:
    markdown = staging_directory / "full.md"
    if not markdown.is_file() or markdown.stat().st_size == 0:
        raise StandardParseError(f"MinerU result is missing full.md: {markdown}")
    outputs = sorted(
        str(path.relative_to(staging_directory)).replace("\\", "/")
        for path in staging_directory.rglob("*")
        if path.is_file()
    )
    if not outputs:
        raise StandardParseError("MinerU result contains no files")
    return outputs


def promote_outputs(
    staging_directory: Path,
    run_directory: Path,
    previous_outputs: list[str],
) -> list[str]:
    outputs = staged_outputs(staging_directory)
    top_level_names = {Path(value).parts[0] for value in outputs}
    reserved_collision = top_level_names & RESERVED_OUTPUT_NAMES
    if reserved_collision:
        names = ", ".join(sorted(reserved_collision))
        raise StandardParseError(f"MinerU result collides with reserved cache files: {names}")

    previous_roots = {Path(value).parts[0] for value in previous_outputs if Path(value).parts}
    for name in previous_roots:
        if name in RESERVED_OUTPUT_NAMES:
            continue
        target = run_directory / name
        if target.is_dir():
            shutil.rmtree(target)
        else:
            target.unlink(missing_ok=True)

    for child in staging_directory.iterdir():
        destination = run_directory / child.name
        if destination.exists():
            raise StandardParseError(
                f"Cannot replace untracked file in standard cache: {destination}"
            )
        shutil.move(str(child), destination)
    staging_directory.rmdir()
    return outputs


def result_payload(
    standard_number: str,
    pdf_path: Path,
    run_directory: Path,
    manifest_path: Path,
    pages: str | None,
    pdf_cached: bool,
    parse_cached: bool,
    service: MineruServiceConfig,
    outputs: list[str],
) -> dict[str, object]:
    markdown = run_directory / "full.md"
    if not markdown.is_file() or markdown.stat().st_size == 0:
        raise StandardParseError(f"MinerU result is missing full.md: {markdown}")
    return {
        "standard_number": standard_number,
        "cached": parse_cached,
        "pdf_cached": pdf_cached,
        "parse_cached": parse_cached,
        "pages": pages,
        "pdf": str(pdf_path.resolve()),
        "markdown": str(markdown.resolve()),
        "output_directory": str(run_directory.resolve()),
        "outputs": outputs,
        "manifest": str(manifest_path.resolve()),
        "service": {
            "mode": service.mode,
            "submit_endpoint": service.submit_endpoint,
            "result_endpoint_template": service.result_endpoint_template,
            "config": str(service.path),
            "token_source": service.token_source,
            "request_timeout_seconds": service.request_timeout,
            "parse_timeout_seconds": service.parse_timeout,
            "poll_interval_seconds": service.poll_interval,
        },
        "config": parse_profile(pages),
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Download an official GB/GB/T/GB/Z PDF and parse it to Markdown with MinerU."
    )
    parser.add_argument("standard_number", help='standard number such as "GB/T 35697-2026"')
    parser.add_argument("--config", type=Path, help="MinerU service JSON configuration")
    parser.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT)
    parser.add_argument("--pages", help='optional smoke-test range such as "7,11"')
    parser.add_argument("--browser-timeout", type=float, default=120.0)
    return parser


def run(args: argparse.Namespace) -> dict[str, object]:
    standard_number = normalize_standard_number(args.standard_number)
    pages = validate_pages(args.pages)
    if args.browser_timeout <= 0:
        raise StandardParseError("Browser timeout must be positive")
    config_setting = args.config or os.environ.get("MINERU_CONFIG") or DEFAULT_CONFIG_PATH
    service = load_service_config(Path(config_setting))

    suffix = f"__pages-{pages.replace(',', '_')}" if pages else ""
    run_directory = (args.output_root.expanduser().resolve() / f"{standard_slug(standard_number)}{suffix}")
    pdf_path = run_directory / "source.pdf"
    manifest_path = run_directory / "cache-manifest.json"

    run_directory.mkdir(parents=True, exist_ok=True)
    lock_path = run_directory / ".parse.lock"
    try:
        lock = lock_path.open("x", encoding="utf-8")
    except FileExistsError as error:
        raise StandardParseError(f"Another parse is already running: {run_directory}") from error

    try:
        with lock:
            lock.write(json.dumps({"pid": os.getpid(), "started": time.time()}))

        manifest = load_cache_manifest(manifest_path)
        source_sha256 = cached_source_sha256(manifest, standard_number, pdf_path)
        pdf_cached = source_sha256 is not None
        if pdf_cached:
            print(f"Reusing downloaded PDF: {pdf_path}", file=sys.stderr)
        else:
            print(f"Downloading {standard_number} from the official site", file=sys.stderr)
            temporary_pdf = run_directory / f".source-{uuid.uuid4().hex}.pdf"
            try:
                browser_result = download_standard(
                    standard_number,
                    temporary_pdf,
                    args.browser_timeout,
                )
                validate_pdf(temporary_pdf)
                source_sha256 = file_sha256(temporary_pdf)
                os.replace(temporary_pdf, pdf_path)
            finally:
                temporary_pdf.unlink(missing_ok=True)
            manifest = {
                "schema_version": MANIFEST_SCHEMA_VERSION,
                "standard_number": standard_number,
                "pages": pages,
                "source": {
                    "path": pdf_path.name,
                    "sha256": source_sha256,
                    "size": pdf_path.stat().st_size,
                    "detail_url": browser_result.get("detail_url"),
                    "suggested_filename": browser_result.get("suggested_filename"),
                    "downloaded_at": time.time(),
                },
                "parse": None,
            }
            write_cache_manifest(manifest_path, manifest)
            print(f"Downloaded PDF: {pdf_path}", file=sys.stderr)

        if source_sha256 is None or manifest is None:
            raise StandardParseError("PDF cache was not initialized correctly")

        profile_sha256 = parse_profile_sha256(pages)
        outputs = cached_parse_outputs(
            manifest,
            run_directory,
            source_sha256,
            profile_sha256,
        )
        if outputs is not None:
            return result_payload(
                standard_number,
                pdf_path,
                run_directory,
                manifest_path,
                pages,
                pdf_cached=True,
                parse_cached=True,
                service=service,
                outputs=outputs,
            )

        if service.mode == "cloud" and not service.token:
            raise StandardParseError("cloud mode requires a token")

        parsed_manifest = manifest.get("parse")
        previous_outputs = (
            [value for value in parsed_manifest.get("outputs", []) if isinstance(value, str)]
            if isinstance(parsed_manifest, dict)
            and isinstance(parsed_manifest.get("outputs"), list)
            else []
        )
        staging_directory = run_directory / f".mineru-staging-{uuid.uuid4().hex}"
        try:
            parse_with_mineru(pdf_path, staging_directory, pages, service)
            outputs = promote_outputs(
                staging_directory,
                run_directory,
                previous_outputs,
            )
        finally:
            if staging_directory.exists():
                shutil.rmtree(staging_directory)

        manifest["parse"] = {
            "source_sha256": source_sha256,
            "profile_sha256": profile_sha256,
            "profile": parse_profile(pages),
            "outputs": outputs,
            "completed_at": time.time(),
        }
        write_cache_manifest(manifest_path, manifest)
        payload = result_payload(
            standard_number,
            pdf_path,
            run_directory,
            manifest_path,
            pages,
            pdf_cached=pdf_cached,
            parse_cached=False,
            service=service,
            outputs=outputs,
        )
        return payload
    finally:
        lock_path.unlink(missing_ok=True)


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        payload = run(args)
    except StandardParseError as error:
        print(f"error: {error}", file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        print("error: interrupted", file=sys.stderr)
        return 130
    print(json.dumps(payload, ensure_ascii=False, indent=2))
    return 0


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