"""Authenticated API Gateway and ShortcutJob compute offload client."""

from __future__ import annotations

import json
import time
import uuid
from dataclasses import dataclass
from typing import Any, Callable, Mapping, Optional
from urllib.parse import urlencode, urlsplit

from .http import UrllibTransport, json_body


class RemoteJobError(RuntimeError):
    """Raised when remote compute submission or execution fails."""


@dataclass(frozen=True)
class RemoteJob:
    """Normalized Supraworx ShortcutJob state."""

    job_id: str
    shortcut_id: str
    status: str
    finished: bool
    result: Any = None
    error_message: str = ""
    raw: Optional[Mapping[str, Any]] = None

    @property
    def succeeded(self) -> bool:
        return self.finished and self.status == "3" and not self.error_message


class ShortcutJobClient:
    """Submit heavy work via API Gateway and poll a ShortcutJob status URL.

    ``submit_mode="json"`` targets a configured API Gateway endpoint.
    ``submit_mode="form"`` matches ``shortcuts/api-shortcut-starter/``.
    """

    def __init__(
        self,
        *,
        submit_url: str,
        status_url: str,
        shortcut_id: str = "",
        shortcut_ids: Optional[Mapping[str, str]] = None,
        api_key: str = "",
        bearer_token: str = "",
        bearer_token_provider: Optional[Callable[..., str]] = None,
        endpoint_url_provider: Optional[Callable[[str], str]] = None,
        device_uid: str = "",
        device_secret: str = "",
        submit_mode: str = "json",
        transport: Optional[UrllibTransport] = None,
        timeout: float = 30.0,
        max_request_bytes: int = 8 * 1024 * 1024,
        allow_insecure_http: bool = False,
        allow_ambiguous_legacy_polling: bool = False,
        extra_headers: Optional[Mapping[str, str]] = None,
    ) -> None:
        if not allow_insecure_http and (
            not submit_url.startswith("https://")
            or not status_url.startswith("https://")
        ):
            raise ValueError("Remote compute URLs must use HTTPS")
        if submit_mode not in {"json", "form"}:
            raise ValueError("submit_mode must be json or form")
        normalized_shortcut_ids = {
            str(job_type).strip(): str(identifier).strip()
            for job_type, identifier in dict(shortcut_ids or {}).items()
            if str(job_type).strip() and str(identifier).strip()
        }
        if not shortcut_id and not normalized_shortcut_ids:
            raise ValueError("at least one shortcut_id is required")
        self.submit_url = submit_url
        self.status_url = status_url
        self.shortcut_id = str(shortcut_id)
        self.shortcut_ids = normalized_shortcut_ids
        self.api_key = api_key
        self.bearer_token = bearer_token
        self.bearer_token_provider = bearer_token_provider
        self.endpoint_url_provider = endpoint_url_provider
        self.device_uid = device_uid
        self.device_secret = device_secret
        self.submit_mode = submit_mode
        self.transport = transport or UrllibTransport()
        self.timeout = float(timeout)
        if max_request_bytes < 1024:
            raise ValueError("max_request_bytes must be at least 1024")
        self.max_request_bytes = int(max_request_bytes)
        self.allow_insecure_http = bool(allow_insecure_http)
        self.allow_ambiguous_legacy_polling = bool(allow_ambiguous_legacy_polling)
        self.extra_headers = dict(extra_headers or {})

    def submit(
        self,
        payload: Mapping[str, Any],
        *,
        idempotency_key: str = "",
    ) -> RemoteJob:
        """Submit one compute job and return its initial state."""
        key = idempotency_key or str(uuid.uuid4())
        job_type = str(payload.get("job_type") or "").strip()
        shortcut_id = self.shortcut_ids.get(job_type, self.shortcut_id)
        if not shortcut_id:
            raise RemoteJobError(
                "No ShortcutJob is configured for job type {!r}".format(job_type)
            )
        submit_url = self._target_url(self.submit_url)
        headers = self._headers(submit_url)
        headers["Idempotency-Key"] = key
        if self.submit_mode == "json":
            headers["Content-Type"] = "application/json"
            request_payload = {
                "schema_version": "labcobot.remote-job.v1",
                "shortcut_id": shortcut_id,
                "idempotency_key": key,
                "device_uid": self.device_uid,
                "input": dict(payload),
            }
            body = json_body(request_payload)
        else:
            headers["Content-Type"] = "application/x-www-form-urlencoded"
            body = urlencode(
                {
                    "shortcut_id": shortcut_id,
                    "parameters": key,
                    "user_input_text": json.dumps(
                        dict(payload), separators=(",", ":"), allow_nan=False
                    ),
                }
            ).encode("ascii")
        if len(body) > self.max_request_bytes:
            raise RemoteJobError("Remote job request exceeds configured maximum size")
        response = self._request_with_auth_retry(
            "POST",
            submit_url,
            headers=headers,
            body=body,
            timeout=self.timeout,
            max_response_bytes=2 * 1024 * 1024,
        )
        if not 200 <= response.status_code < 300:
            raise RemoteJobError(
                "Remote job submission returned HTTP {}".format(response.status_code)
            )
        data = response.json()
        if not isinstance(data, Mapping):
            raise RemoteJobError("Remote job submission returned invalid JSON")
        if data.get("success") is False:
            raise RemoteJobError(str(data.get("error") or "Remote job was rejected"))
        job = self._normalize(data, fallback_shortcut_id=shortcut_id)
        if not job.job_id and not self.allow_ambiguous_legacy_polling:
            raise RemoteJobError(
                "Remote job submission did not return a correlation job id"
            )
        return job

    def poll(self, job: RemoteJob) -> RemoteJob:
        """Fetch the current state of a submitted ShortcutJob."""
        if not job.job_id and not self.allow_ambiguous_legacy_polling:
            raise RemoteJobError(
                "Cannot poll a remote job without a correlation job id"
            )
        form = {"shortcutid": job.shortcut_id or self.shortcut_id}
        if job.job_id:
            form["shortcutjob_id"] = job.job_id
        else:
            form["latest_running"] = "true"
        status_url = self._target_url(self.status_url)
        response = self._request_with_auth_retry(
            "POST",
            status_url,
            headers={
                **self._headers(status_url),
                "Content-Type": "application/x-www-form-urlencoded",
            },
            body=urlencode(form).encode("ascii"),
            timeout=self.timeout,
            max_response_bytes=2 * 1024 * 1024,
        )
        if not 200 <= response.status_code < 300:
            raise RemoteJobError(
                "Remote job status returned HTTP {}".format(response.status_code)
            )
        data = response.json()
        if not isinstance(data, Mapping):
            raise RemoteJobError("Remote job status returned invalid JSON")
        if data.get("success") is False and job.job_id:
            raise RemoteJobError(str(data.get("error") or "Remote job was not found"))
        if not data.get("id") and not data.get("shortcutjob_id") and not job.job_id:
            # The starter API omits the id. Once no running job remains, query
            # the latest completed result for this shortcut.
            result_form = urlencode(
                {
                    "shortcutid": job.shortcut_id or self.shortcut_id,
                    "latest_result": "true",
                }
            ).encode("ascii")
            result_response = self._request_with_auth_retry(
                "POST",
                status_url,
                headers={
                    **self._headers(status_url),
                    "Content-Type": "application/x-www-form-urlencoded",
                },
                body=result_form,
                timeout=self.timeout,
                max_response_bytes=2 * 1024 * 1024,
            )
            if 200 <= result_response.status_code < 300:
                result_data = result_response.json()
                if isinstance(result_data, Mapping):
                    data = result_data
        self._require_matching_job(data, job)
        normalized = self._normalize(
            data,
            fallback_job_id=job.job_id,
            fallback_shortcut_id=job.shortcut_id or self.shortcut_id,
        )
        return normalized

    @staticmethod
    def _require_matching_job(
        data: Mapping[str, Any],
        expected: RemoteJob,
    ) -> None:
        """Reject status payloads correlated to a different job or shortcut."""

        returned_job_id = str(
            data.get("shortcutjob_id") or data.get("job_id") or data.get("id") or ""
        )
        if expected.job_id and returned_job_id and returned_job_id != expected.job_id:
            raise RemoteJobError("Remote job status returned a different job id")

        returned_shortcut_id = str(data.get("shortcut_id") or "")
        expected_shortcut_id = expected.shortcut_id or ""
        if (
            expected_shortcut_id
            and returned_shortcut_id
            and returned_shortcut_id != expected_shortcut_id
        ):
            raise RemoteJobError("Remote job status returned a different shortcut id")

    def wait(
        self,
        job: RemoteJob,
        *,
        deadline_seconds: float = 300.0,
        poll_interval_seconds: float = 2.0,
    ) -> RemoteJob:
        """Poll until completion or a bounded deadline."""
        deadline = time.monotonic() + max(0.0, deadline_seconds)
        current = job
        while not current.finished:
            if time.monotonic() >= deadline:
                raise TimeoutError("Remote ShortcutJob did not finish before deadline")
            time.sleep(max(0.05, poll_interval_seconds))
            current = self.poll(current)
        if not current.succeeded:
            raise RemoteJobError(
                current.error_message
                or (
                    "Remote ShortcutJob finished without the required "
                    "successful status"
                )
            )
        return current

    def _headers(self, target_url: str, *, force_refresh: bool = False) -> dict:
        headers = {
            "Accept": "application/json",
            "User-Agent": "LabCobot-Offload/1",
            **self.extra_headers,
        }
        if self.api_key:
            headers["X-API-Key"] = self.api_key
        bearer_token = ""
        if self.bearer_token_provider is not None:
            try:
                bearer_token = str(
                    self.bearer_token_provider(
                        target_url,
                        force_refresh=force_refresh,
                    )
                    or ""
                ).strip()
            except Exception as exc:
                raise RemoteJobError(
                    "Tenant login is unavailable for the remote endpoint"
                ) from exc
        elif self.bearer_token:
            bearer_token = self.bearer_token
        if bearer_token:
            headers["Authorization"] = "Bearer {}".format(bearer_token)
        if self.device_uid:
            headers["X-Device-Id"] = self.device_uid
        if self.device_secret:
            headers["X-Device-Secret"] = self.device_secret
        return headers

    def _target_url(self, configured_url: str) -> str:
        """Resolve a commissioned path against the authenticated tenant."""

        if self.endpoint_url_provider is None:
            return configured_url
        try:
            resolved = str(self.endpoint_url_provider(configured_url) or "").strip()
        except Exception as exc:
            raise RemoteJobError(
                "Tenant endpoint is unavailable for the remote request"
            ) from exc
        if not resolved:
            raise RemoteJobError(
                "Tenant endpoint is unavailable for the remote request"
            )
        configured = urlsplit(configured_url)
        target = urlsplit(resolved)
        try:
            target_port = target.port or (
                443 if target.scheme.lower() == "https" else 80
            )
        except ValueError as exc:
            raise RemoteJobError(
                "Tenant endpoint is unavailable for the remote request"
            ) from exc
        secure_target = target.scheme.lower() == "https" and target_port == 443
        local_development_target = bool(
            self.allow_insecure_http
            and target.scheme.lower() == "http"
            and target_port == 80
            and target.hostname in {"localhost", "127.0.0.1", "::1"}
        )
        if (
            target.path != configured.path
            or target.query != configured.query
            or target.fragment
            or not target.hostname
            or target.username is not None
            or target.password is not None
            or not (secure_target or local_development_target)
        ):
            raise RemoteJobError(
                "Tenant endpoint is unavailable for the remote request"
            )
        return resolved

    def _request_with_auth_retry(
        self,
        method: str,
        url: str,
        *,
        headers: Mapping[str, str],
        body: bytes,
        timeout: float,
        max_response_bytes: int,
    ):
        """Retry one 401 once after a tenant-bound refresh."""

        response = self.transport.request(
            method,
            url,
            headers=headers,
            body=body,
            timeout=timeout,
            max_response_bytes=max_response_bytes,
        )
        if response.status_code != 401 or self.bearer_token_provider is None:
            return response
        refreshed_headers = dict(headers)
        refreshed_headers.pop("Authorization", None)
        current_headers = self._headers(url, force_refresh=True)
        authorization = current_headers.get("Authorization")
        if not authorization:
            return response
        refreshed_headers["Authorization"] = authorization
        return self.transport.request(
            method,
            url,
            headers=refreshed_headers,
            body=body,
            timeout=timeout,
            max_response_bytes=max_response_bytes,
        )

    @staticmethod
    def _normalize(
        data: Mapping[str, Any],
        *,
        fallback_job_id: str = "",
        fallback_shortcut_id: str = "",
    ) -> RemoteJob:
        job_id = str(
            data.get("shortcutjob_id")
            or data.get("job_id")
            or data.get("id")
            or fallback_job_id
            or ""
        )
        shortcut_id = str(data.get("shortcut_id") or fallback_shortcut_id or "")
        status = str(data.get("status") or data.get("task_status") or "1")
        if status not in {"0", "1", "2", "3", "4"}:
            raise RemoteJobError("Remote job returned an unknown status")
        finished = bool(data.get("finished")) or status in {"3", "4"}
        result = data.get("shortcutjob_result", data.get("result"))
        error = str(data.get("error_message") or data.get("error") or "")
        return RemoteJob(
            job_id=job_id,
            shortcut_id=shortcut_id,
            status=status,
            finished=finished,
            result=result,
            error_message=error,
            raw=dict(data),
        )
