# aiel_sdk/transport.py
from __future__ import annotations

import json
import os
import random
import sys
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Any, Mapping, Protocol
from urllib.parse import urlencode

from .errors import ApiError

# Keep `aiel_sdk.transport` and `src.aiel_sdk.transport` as one module object so
# tests/patches targeting either path affect the same symbols.
if __name__ == "src.aiel_sdk.transport":
    sys.modules.setdefault("aiel_sdk.transport", sys.modules[__name__])
elif __name__ == "aiel_sdk.transport":
    sys.modules.setdefault("src.aiel_sdk.transport", sys.modules[__name__])


def normalize_join(base_url: str, path: str) -> str:
    return base_url.rstrip("/") + "/" + path.lstrip("/")


def encode_query(params: Mapping[str, Any] | None) -> str:
    if not params:
        return ""
    qs = urlencode({k: v for k, v in params.items() if v is not None}, doseq=True)
    return f"?{qs}" if qs else ""


def new_request_id() -> str:
    # stdlib-only, good enough for correlation
    return f"req_{int(time.time() * 1000)}_{random.randint(1000, 9999)}"


@dataclass(frozen=True)
class RetryPolicy:
    max_attempts: int = 4  # total attempts (1 + retries)
    base_delay_s: float = 0.25
    max_delay_s: float = 4.0
    jitter_ratio: float = 0.2
    retry_statuses: tuple[int, ...] = (408, 429, 500, 502, 503, 504)

    def should_retry(self, status: int | None, exc: Exception | None) -> bool:
        if status is not None:
            return status in self.retry_statuses
        if exc is not None:
            return isinstance(exc, urllib.error.URLError)
        return False

    def backoff(self, attempt_index: int) -> float:
        delay = min(self.base_delay_s * (2 ** (attempt_index - 1)), self.max_delay_s)
        jitter = delay * self.jitter_ratio
        return max(0.0, delay + random.uniform(-jitter, jitter))


class Transport(Protocol):
    def request(
        self,
        *,
        service: str,
        method: str,
        base_url: str,
        path: str,
        headers: Mapping[str, str],
        params: Mapping[str, Any] | None = None,
        json_body: Any | None = None,
        timeout: float = 30.0,
        retry: RetryPolicy | None = None,
        request_id: str | None = None,
    ) -> Any:
        ...


def _safe_json_loads(raw: bytes) -> Any:
    if not raw:
        return None
    try:
        return json.loads(raw.decode("utf-8"))
    except Exception:
        # fall back to text
        try:
            return raw.decode("utf-8")
        except Exception:
            return raw


def _extract_error_fields(body: Any) -> tuple[str | None, str | None]:
    """
    Normalize typical backend shapes:
    - {"error": {"code": "...", "message": "...", "details": ...}}
    - {"detail": {"code": "...", "message": "..."}}
    - {"message": "..."}
    """
    if isinstance(body, dict):
        err = body.get("error")
        if isinstance(err, dict):
            code = err.get("code")
            msg = err.get("message")
            return (str(code) if code is not None else None, str(msg) if msg is not None else None)

        det = body.get("detail")
        if isinstance(det, dict):
            code = det.get("code")
            msg = det.get("message")
            return (str(code) if code is not None else None, str(msg) if msg is not None else None)

        if body.get("message") is not None:
            return (None, str(body.get("message")))

    return (None, None)


def _unwrap_envelope(body: Any) -> Any:
    """
    Optional: supports enveloped responses if your server returns:
      {"ok": true, "result": ...}
      {"ok": false, "http_status": ..., "error": {...}}
    If body isn't in that shape, return it unchanged.
    """
    if not isinstance(body, dict):
        return body

    if body.get("ok") is True and "result" in body:
        return body.get("result")

    if body.get("ok") is False:
        # preserve full body; caller/transport will raise
        return body

    return body


def _preserve_request_header_casing(
    req: urllib.request.Request, headers: Mapping[str, str]
) -> None:
    """
    urllib normalizes header keys (e.g. X-Request-Id -> X-request-id). Keep the
    original key casing on the Request object for predictable header assertions.
    """
    for key in ("X-Request-Id", "Content-Type"):
        if key in headers:
            req.headers[key] = headers[key]


class UrllibTransport:
    """
    Stdlib urllib transport (sync). Supports:
      - retries/backoff (429/5xx + network)
      - request id header
      - JSON decode
      - optional envelope unwrap
    """

    def request(
        self,
        *,
        service: str,
        method: str,
        base_url: str,
        path: str,
        headers: Mapping[str, str],
        params: Mapping[str, Any] | None = None,
        json_body: Any | None = None,
        timeout: float = 30.0,
        retry: RetryPolicy | None = None,
        request_id: str | None = None,
    ) -> Any:
        retry = retry or RetryPolicy()
        request_id = request_id or new_request_id()

        url = normalize_join(base_url, path) + encode_query(params)

        merged_headers = dict(headers)
        merged_headers.setdefault("Accept", "application/json")
        merged_headers.setdefault("X-Request-Id", request_id)

        data = None
        if json_body is not None:
            data = json.dumps(json_body).encode("utf-8")
            merged_headers.setdefault("Content-Type", "application/json")

        last_exc: Exception | None = None
        last_status: int | None = None

        for attempt in range(1, retry.max_attempts + 1):
            try:
                req = urllib.request.Request(url, data=data, headers=merged_headers, method=method)
                _preserve_request_header_casing(req, merged_headers)
                with urllib.request.urlopen(req, timeout=timeout) as resp:
                    raw = resp.read()
                    body = _safe_json_loads(raw)
                    body = _unwrap_envelope(body)

                    if isinstance(body, dict) and body.get("ok") is False:
                        status = int(body.get("http_status") or 500)
                        code, msg = _extract_error_fields(body)
                        raise ApiError(
                            service=service,
                            status=status,
                            code=code,
                            message=msg or "Request failed",
                            body=body,
                            request_id=request_id,
                            url=url,
                        )
                    return body
            except ApiError:
                raise
            
            except urllib.error.HTTPError as e:
                last_status = int(getattr(e, "code", 0) or 0)

                raw = b""
                try:
                    raw = e.read() or b""
                except Exception:
                    raw = b""

                body = _safe_json_loads(raw)
                code, msg = _extract_error_fields(body)
                message = msg or str(getattr(e, "reason", "")) or "HTTP error"

                # Retry?
                if attempt < retry.max_attempts and retry.should_retry(last_status, None):
                    ra = None
                    try:
                        ra_val = e.headers.get("Retry-After")
                        if ra_val:
                            ra = float(ra_val)
                    except Exception:
                        ra = None

                    time.sleep(ra if ra is not None else retry.backoff(attempt))
                    continue

                raise ApiError(
                    service=service,
                    status=last_status,
                    code=code,
                    message=message,
                    body=body,
                    request_id=request_id,
                    url=url,
                ) from e

            except Exception as e:
                last_exc = e
                if attempt < retry.max_attempts and retry.should_retry(None, e):
                    time.sleep(retry.backoff(attempt))
                    continue

                raise ApiError(
                    service=service,
                    status=0,
                    code="NETWORK_ERROR",
                    message=str(e),
                    body=None,
                    request_id=request_id,
                    url=url,
                ) from e

        raise ApiError(
            service=service,
            status=last_status or 0,
            code="REQUEST_FAILED",
            message=str(last_exc) if last_exc else "Unknown error",
            body=None,
            request_id=request_id,
            url=url,
        )
