#!/usr/bin/env python3
"""Wait for an Atomelody task; emit JSON Lines only on changes or exit.

Requires Python 3.9+ and ATOMELODY_API_KEY. Exit codes: 0 completed,
1 failed/cancelled, 2 query/configuration error, 3 wait timeout, 130 interrupted.
"""

import argparse
import json
import math
import os
import random
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timezone
from email.utils import parsedate_to_datetime


class QueryError(Exception):
    def __init__(self, message, *, retryable=False, retry_after=0):
        super().__init__(message)
        self.retryable = retryable
        self.retry_after = retry_after


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        # Never forward the API Key to a redirected destination.
        return None


def retry_delay(value):
    try:
        seconds = float(value)
    except (ValueError, TypeError):
        try:
            seconds = parsedate_to_datetime(value).timestamp() - time.time()
        except (ValueError, TypeError, OverflowError):
            return 0
    return max(0, seconds) if math.isfinite(seconds) else 0


def query_status(opener, request, timeout):
    try:
        with opener.open(request, timeout=timeout) as response:
            body = json.load(response)
    except urllib.error.HTTPError as exc:
        raise QueryError(
            f"HTTP {exc.code}",
            retryable=exc.code == 429 or 500 <= exc.code < 600,
            retry_after=retry_delay(exc.headers.get("Retry-After")),
        ) from exc
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        raise QueryError("Network request failed", retryable=True) from exc
    except (ValueError, UnicodeError) as exc:
        raise QueryError("Invalid JSON response") from exc

    if not isinstance(body, dict) or body.get("code") != "OK":
        raise QueryError("API response was not OK")
    task = body.get("data")
    if not isinstance(task, dict) or type(task.get("status")) is not int or task["status"] not in range(1, 7):
        raise QueryError("Invalid task status response")
    if not isinstance(task.get("task_id"), str):
        raise QueryError("Missing task ID")
    if "output_url" in task and not isinstance(task["output_url"], str):
        raise QueryError("Invalid result URL")
    if task.get("error") is not None:
        if not isinstance(task["error"], dict) or not isinstance(task["error"].get("code"), str):
            raise QueryError("Invalid task error")
    return task


def emit(event, task_id, **fields):
    print(json.dumps({
        "event": event,
        "task_id": task_id,
        "observed_at": datetime.now(timezone.utc).isoformat(),
        **fields,
    }, ensure_ascii=False), flush=True)


def wait_task(task_id, api_key, max_wait):
    request = urllib.request.Request(
        "https://api.atomelody.com/v1/tasks/status?" + urllib.parse.urlencode({"task_id": task_id}),
        headers={"Authorization": f"Bearer {api_key}", "Accept": "application/json"},
    )
    opener = urllib.request.build_opener(NoRedirect())
    deadline = time.monotonic() + max_wait
    last_status = None
    interval = 15
    errors = 0

    while time.monotonic() < deadline:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            break
        server_delay = 0
        try:
            task = query_status(opener, request, min(30, remaining))
            if task["task_id"] != task_id:
                raise QueryError("Response task ID does not match")
            status = task["status"]
            errors = 0
            if status in (4, 5, 6):
                result = {"status": status}
                if task.get("output_url"):
                    result["output_url"] = task["output_url"]
                if task.get("error"):
                    result["error"] = {"code": task["error"]["code"]}
                emit({4: "completed", 5: "failed", 6: "cancelled"}[status], task_id, **result)
                return 0 if status == 4 else 1
            if status != last_status:
                emit("status_changed", task_id, previous_status=last_status, status=status)
                last_status = status
        except QueryError as exc:
            errors += 1
            if not exc.retryable or errors >= 5:
                emit("error", task_id, message=str(exc), consecutive_errors=errors)
                return 2
            server_delay = exc.retry_after

        remaining = deadline - time.monotonic()
        if remaining > 0:
            # Jitter avoids synchronized clients. A server Retry-After takes precedence.
            delay = max(min(120, interval * random.uniform(0.8, 1.2)), server_delay)
            time.sleep(min(delay, remaining))
        interval = min(interval * 2, 120)

    emit("pending", task_id, status=last_status, reason="wait_timeout")
    return 3


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("task_id")
    parser.add_argument("--max-wait", type=float, default=1800, help="Wait budget in seconds (default: 1800)")
    args = parser.parse_args()
    api_key = os.environ.get("ATOMELODY_API_KEY", "").strip()
    if not api_key or not args.task_id.strip() or not math.isfinite(args.max_wait) or args.max_wait <= 0:
        emit("error", args.task_id, message="Set ATOMELODY_API_KEY and use a positive finite max-wait and nonempty task_id")
        return 2
    try:
        return wait_task(args.task_id, api_key, args.max_wait)
    except KeyboardInterrupt:
        emit("interrupted", args.task_id, message="Local waiting stopped; remote task was not cancelled")
        return 130


if __name__ == "__main__":
    sys.exit(main())