"""Standard-library paid API example. Read credentials only from VTN_API_KEY."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from contextlib import ExitStack
from pathlib import Path

TERMINAL = {"completed", "completed_with_gaps", "interrupted", "cancelled", "failed", "expired"}


class ApiError(Exception):
    pass


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None  # Never forward service credentials or follow an upload redirect.


class Client:
    def __init__(self, base, key, *, opener=None, sleep=time.sleep):
        parsed = urllib.parse.urlsplit(base)
        if (
            parsed.scheme != "https"
            or not parsed.netloc
            or parsed.query
            or parsed.fragment
            or parsed.username is not None
            or parsed.password is not None
        ):
            raise ApiError(
                "VTN_API_BASE must be an HTTPS API origin, including any deployment prefix."
            )
        if not key or "\n" in key or "\r" in key:
            raise ApiError("Set VTN_API_KEY in your server environment.")
        self.base, self.key = base.rstrip("/"), key
        self.opener = opener or urllib.request.build_opener(NoRedirect())
        self.sleep = sleep

    def request(self, method, path, payload=None, *, idempotency=None):
        if not path.startswith("/") or path.startswith("//"):
            raise ApiError("Use an API-relative resource path.")
        headers = {"Authorization": "Bearer " + self.key, "Accept": "application/json"}
        if payload is not None:
            headers["Content-Type"] = "application/json"
        if idempotency:
            headers["Idempotency-Key"] = idempotency
        body = json.dumps(payload).encode() if payload is not None else None
        for attempt in range(6):
            request = urllib.request.Request(
                self.base + path, data=body, headers=headers, method=method
            )
            try:
                with self.opener.open(request, timeout=120) as response:
                    return json.load(response)
            except urllib.error.HTTPError as error:
                # Only the safe error code is reported, never headers, keys or URLs.
                if error.code == 429 and attempt < 5:
                    try:
                        delay = max(1, int(error.headers.get("Retry-After", "60")))
                    except ValueError:
                        delay = 60
                    error.close()
                    self.sleep(delay)
                    continue
                try:
                    code = json.load(error).get("error", {}).get("code", "HTTP_ERROR")
                except (ValueError, AttributeError):
                    code = "HTTP_ERROR"
                raise ApiError(f"API request failed: HTTP {error.code} {code}.") from None
            except (urllib.error.URLError, TimeoutError):
                # A timed-out job POST is safe to repeat with its saved identity.
                raise ApiError(
                    "Network request failed. Retry using the same saved request state."
                ) from None
        raise ApiError("Request retry limit reached.")

    def unsigned(self, method, url, data=None):
        if urllib.parse.urlsplit(url).scheme != "https":
            raise ApiError("The service returned an invalid storage URL.")
        # Storage requests deliberately contain no API Authorization header.
        request = urllib.request.Request(url, data=data, method=method)
        try:
            return self.opener.open(request, timeout=120)
        except (urllib.error.URLError, TimeoutError):
            raise ApiError(
                "Storage transfer failed. Request a fresh URL or retry the upload."
            ) from None

    def upload(self, source):
        size = source.stat().st_size
        capabilities = self.request("GET", "/api/v1/account")
        limits = capabilities["capabilities"]["limits"]
        if (
            not capabilities["eligible"]
            or not capabilities["capabilities"]["api_admission_enabled"]
        ):
            raise ApiError("API admission is unavailable for this account.")
        extension = source.suffix.lower().lstrip(".")
        if (
            extension not in capabilities["capabilities"]["formats"]
            or not 0 < size <= limits["max_upload_bytes"]
        ):
            raise ApiError("Choose a supported, nonempty file within the advertised byte limit.")
        upload = self.request(
            "POST", "/api/v1/uploads", {"byte_size": size, "extension": extension}
        )
        path = "/api/v1/uploads/" + urllib.parse.quote(upload["upload_id"], safe="")
        completed = []
        try:
            with source.open("rb") as media:
                for number in range(1, upload["part_count"] + 1):
                    targets = self.request("POST", path + "/parts", {"part_numbers": [number]})
                    chunk = media.read(upload["part_size"])
                    with self.unsigned("PUT", targets["targets"][0]["url"], chunk) as response:
                        etag = response.headers.get("ETag")
                    if not etag:
                        raise ApiError("Storage did not return the required ETag.")
                    completed.append({"part_number": number, "etag": etag})  # Keep quotes.
            self.request("POST", path + "/complete", {"parts": completed})
            return upload["upload_id"]
        except BaseException:
            try:
                self.request("POST", path + "/abort", {})
            except ApiError:
                pass  # The normal owned-upload expiry remains in effect.
            raise

    def wait(self, job, *, timeout=7200):
        deadline = time.monotonic() + timeout
        while job["state"] not in TERMINAL:
            if time.monotonic() >= deadline:
                raise ApiError(
                    "Polling deadline reached. Use status/result or cancel with the saved job ID."
                )
            self.sleep(max(5, job.get("poll_after_seconds", 5)))
            job = self.request("GET", "/api/v1/jobs/" + urllib.parse.quote(job["job_id"], safe=""))
        return job

    def results(self, job_id):
        cursor = None
        while True:
            path = "/api/v1/jobs/" + urllib.parse.quote(job_id, safe="") + "/result"
            if cursor:
                path += "?" + urllib.parse.urlencode({"cursor": cursor})
            page = self.request("GET", path)
            yield page
            cursor = page["next_cursor"]
            if cursor is None:
                break


def save_json(path, payload):
    # State contains only upload/job/request identity, never credentials or signed URLs.
    temporary = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="w", encoding="utf-8", dir=path.parent, prefix=".api-request-", delete=False
        ) as stream:
            temporary = Path(stream.name)
            stream.write(json.dumps(payload, indent=2) + "\n")
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
    finally:
        if temporary is not None:
            temporary.unlink(missing_ok=True)


def run(client, args):
    source = args.file.resolve()
    options = {}
    for name in ("audio_language", "notes_language"):
        if getattr(args, name, None):
            options[name] = getattr(args, name)
    start, end = getattr(args, "start", None), getattr(args, "end", None)
    if start is not None or end is not None:
        if (
            start is None
            or end is None
            or not math.isfinite(start)
            or not math.isfinite(end)
            or not 0 <= start < end
        ):
            raise ApiError(
                "Range requires finite --start and --end seconds, with 0 <= start < end."
            )
        options["range"] = {"start_seconds": start, "end_seconds": end}
    with source.open("rb") as stream:
        fingerprint = hashlib.file_digest(stream, "sha256").hexdigest()
    if args.state.exists():
        saved = json.loads(args.state.read_text(encoding="utf-8"))
        if saved["sha256"] != fingerprint or saved["api_base"] != client.base:
            raise ApiError(
                "Saved state belongs to a different file or environment. Choose a new state path."
            )
        if saved.get("options", {}) != options:
            raise ApiError(
                "Saved request has different options. Replay it with the original options."
            )
    else:
        saved = {
            "api_base": client.base,
            "sha256": fingerprint,
            "upload_id": client.upload(source),
            "idempotency_key": uuid.uuid4().hex,
            "options": options,
        }
        save_json(args.state, saved)  # Persist BEFORE job creation so lost responses are safe.
    payload = {"upload_id": saved["upload_id"], **options}
    job = client.request("POST", "/api/v1/jobs", payload, idempotency=saved["idempotency_key"])
    saved["job_id"] = job["job_id"]
    save_json(args.state, saved)
    job = client.wait(job, timeout=args.timeout)
    print(f"Job {job['job_id']}: {job['state']}, charged seconds: {job['charged_seconds']}")
    if job["result_available"]:
        export_results(client, job["job_id"], args.output, args.download)
    elif job["state"] != "completed":
        raise ApiError(
            "Job has no retained result. Inspect its status error before starting another job."
        )


def export_results(client, job_id, output, download=None):
    # Stream pages to disk rather than building an unbounded result list in memory.
    if download:
        download.mkdir(parents=True, exist_ok=True)
    artifact_index = 0
    with ExitStack() as files, output.open("w", encoding="utf-8") as stream:
        partial_files = {}
        for page in client.results(job_id):
            # JSONL preserves partial results, gaps, timing and exact charged usage.
            stream.write(json.dumps(page, ensure_ascii=False) + "\n")
            if download and page.get("partial"):
                for item in page["items"]:
                    kind = item["kind"]
                    if kind not in {"transcript", "notes", "gap"}:
                        continue
                    name = "transcript" if kind == "gap" else kind
                    if name not in partial_files:
                        suffix = "txt" if name == "transcript" else "md"
                        partial_files[name] = files.enter_context(
                            (download / f"partial-{name}.{suffix}").open("w", encoding="utf-8")
                        )
                        partial_files[name].write(
                            "Partial results. Some transcript or notes sections are incomplete.\n\n"
                        )
                    text = (
                        item.get("text", "")
                        if kind != "gap"
                        else (
                            f"[missing: {item['gap']['start_seconds']}–"
                            f"{item['gap']['end_seconds']} seconds]"
                        )
                    )
                    partial_files[name].write(text + "\n\n")
            if download:
                for artifact in page["artifacts"]:
                    suffix = {
                        "transcript_txt": "txt",
                        "notes_md": "md",
                        "notes_html": "html",
                        "notes_pdf": "pdf",
                    }.get(artifact["format"])
                    if suffix is None:
                        raise ApiError("Unknown download format.")
                    artifact_index += 1
                    destination = download / f"artifact-{artifact_index:04d}.{suffix}"
                    with (
                        client.unsigned("GET", artifact["url"]) as response,
                        destination.open("wb") as target,
                    ):
                        while chunk := response.read(1024 * 1024):
                            target.write(chunk)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    commands.add_parser("account")
    create = commands.add_parser("run")
    create.add_argument("file", type=Path)
    create.add_argument("--state", type=Path, required=True)
    create.add_argument("--output", type=Path, default=Path("results.jsonl"))
    create.add_argument("--download", type=Path)
    create.add_argument("--timeout", type=int, default=7200)
    create.add_argument("--audio-language")
    create.add_argument("--notes-language")
    create.add_argument("--start", type=float)
    create.add_argument("--end", type=float)
    for command in ("status", "cancel", "resume", "result"):
        sub = commands.add_parser(command)
        sub.add_argument("job_id")
        if command == "result":
            sub.add_argument("--output", type=Path, default=Path("results.jsonl"))
            sub.add_argument("--download", type=Path)
    args = parser.parse_args()
    try:
        client = Client(os.getenv("VTN_API_BASE", ""), os.getenv("VTN_API_KEY", ""))
        if args.command == "run":
            run(client, args)
        elif args.command == "account":
            print(json.dumps(client.request("GET", "/api/v1/account"), indent=2))
        elif args.command == "result":
            export_results(client, args.job_id, args.output, args.download)
        else:
            path = "/api/v1/jobs/" + urllib.parse.quote(args.job_id, safe="")
            payload = (
                client.request("GET", path)
                if args.command == "status"
                else client.request("POST", path + "/" + args.command, {})
            )
            print(json.dumps(payload, indent=2))
    except (ApiError, OSError, ValueError, KeyError) as error:
        # OSError may contain local paths, but never dump remote requests or headers.
        print(
            str(error)
            if isinstance(error, ApiError)
            else "Local input/state could not be read or written."
        )
        return 1
    return 0


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