import argparse
import os
import time
from pathlib import Path

import requests


def submit_task(base_url: str, token: str, file_path: Path, backend: str) -> str:
    with file_path.open("rb") as handle:
        response = requests.post(
            f"{base_url.rstrip('/')}/tasks",
            headers={"Authorization": f"Bearer {token}"},
            files={"files": (file_path.name, handle, "application/octet-stream")},
            data={"backend": backend, "return_md": "true", "return_images": "true"},
            timeout=300,
        )
    response.raise_for_status()
    payload = response.json()
    task_id = payload.get("task_id") or payload.get("id")
    if not task_id:
        raise RuntimeError(f"No task_id in response: {payload}")
    return task_id


def wait_for_result(base_url: str, token: str, task_id: str, output_path: Path) -> None:
    headers = {"Authorization": f"Bearer {token}"}
    while True:
        status = requests.get(
            f"{base_url.rstrip('/')}/tasks/{task_id}",
            headers=headers,
            timeout=60,
        )
        status.raise_for_status()
        payload = status.json()
        state = str(payload.get("state") or payload.get("status") or "").lower()
        if state in {"done", "completed", "success", "succeeded", "finished"}:
            break
        if state in {"failed", "error"}:
            raise RuntimeError(f"MinerU task failed: {payload}")
        time.sleep(5)

    result = requests.get(
        f"{base_url.rstrip('/')}/tasks/{task_id}/result",
        headers=headers,
        timeout=300,
    )
    result.raise_for_status()
    output_path.write_bytes(result.content)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("file", type=Path)
    parser.add_argument("--base-url", default=os.environ.get("SPACE_URL"))
    parser.add_argument("--token", default=os.environ.get("API_TOKEN"))
    parser.add_argument("--backend", default=os.environ.get("MINERU_BACKEND", "pipeline"))
    parser.add_argument("--output", type=Path, default=Path("mineru-result.bin"))
    args = parser.parse_args()

    if not args.base_url:
        raise SystemExit("Provide --base-url or set SPACE_URL.")
    if not args.token:
        raise SystemExit("Provide --token or set API_TOKEN.")

    task_id = submit_task(args.base_url, args.token, args.file, args.backend)
    print(f"Submitted task {task_id}")
    wait_for_result(args.base_url, args.token, task_id, args.output)
    print(f"Saved {args.output}")


if __name__ == "__main__":
    main()