from __future__ import annotations

import json
import mimetypes
import os
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from urllib import error as urllib_error
from urllib import request as urllib_request
from urllib.parse import urlparse

from .errors import (
    HardsimAuthenticationError,
    HardsimConfigurationError,
    HardsimNotFoundError,
    HardsimRateLimitError,
    HardsimRequestError,
    HardsimServerError,
    HardsimTransportError,
)
from .job import Job

DEFAULT_API_URL = "http://localhost:8000"
DEFAULT_TIMEOUT_S = 30.0
DEFAULT_MAX_RETRIES = 3
DEFAULT_RETRY_BACKOFF_S = 0.5
RETRYABLE_STATUS_CODES = {429, 500, 502, 503, 504}
RETRYABLE_METHODS = {"GET"}


@dataclass
class _SimpleResponse:
    status_code: int
    content: bytes
    headers: dict[str, str]

    @property
    def text(self) -> str:
        return self.content.decode("utf-8", errors="replace")

    def json(self) -> Any:
        return json.loads(self.text)


class _UrlLibSession:
    def __init__(self) -> None:
        self.headers: dict[str, str] = {}

    def request(
        self,
        *,
        method: str,
        url: str,
        json: dict[str, Any] | None = None,
        headers: dict[str, str] | None = None,
        include_default_headers: bool = True,
        timeout: float | None = None,
    ) -> _SimpleResponse:
        merged_headers = dict(self.headers) if include_default_headers else {}
        if headers:
            merged_headers.update(headers)

        payload: bytes | None = None
        if json is not None:
            payload = json_module_dumps(json).encode("utf-8")
            merged_headers.setdefault("Content-Type", "application/json")

        request = urllib_request.Request(
            url=url,
            method=method.upper(),
            headers=merged_headers,
            data=payload,
        )
        request_timeout = 30.0 if timeout is None else float(timeout)
        with urllib_request.urlopen(request, timeout=request_timeout) as response:
            status_code = getattr(response, "status", None) or response.getcode()
            raw_headers = response.headers
            normalized_headers = {str(k): str(v) for k, v in raw_headers.items()}
            body = response.read()
            return _SimpleResponse(
                status_code=int(status_code),
                content=body,
                headers=normalized_headers,
            )


def json_module_dumps(payload: Any) -> str:
    return json.dumps(payload, ensure_ascii=True, separators=(",", ":"))


def _to_response_from_http_error(exc: urllib_error.HTTPError) -> _SimpleResponse:
    body = b""
    try:
        body = exc.read() or b""
    except Exception:
        body = b""
    return _SimpleResponse(status_code=int(exc.code), content=body, headers={})


def _error_detail(response: Any) -> str:
    try:
        payload = response.json()
    except ValueError:
        text = response.text.strip()
        return text or "request failed"

    if isinstance(payload, dict) and "detail" in payload:
        detail = payload.get("detail")
        if isinstance(detail, str):
            return detail
        return json.dumps(detail, ensure_ascii=True)
    return json.dumps(payload, ensure_ascii=True)


def _raise_http_error(response: Any) -> None:
    status = response.status_code
    detail = _error_detail(response)
    if status in {401, 403}:
        raise HardsimAuthenticationError(status_code=status, detail=detail)
    if status == 404:
        raise HardsimNotFoundError(status_code=status, detail=detail)
    if status == 429:
        raise HardsimRateLimitError(status_code=status, detail=detail)
    if status >= 500:
        raise HardsimServerError(status_code=status, detail=detail)
    raise HardsimRequestError(status_code=status, detail=detail)


def _is_absolute_url(value: str) -> bool:
    parsed = urlparse(value)
    return bool(parsed.scheme and parsed.netloc)


def _normalize_optional_idempotency_key(idempotency_key: str | None) -> str | None:
    if idempotency_key is None:
        return None
    key = idempotency_key.strip()
    if not key:
        raise HardsimConfigurationError("idempotency_key must be non-empty when provided")
    if len(key) > 256:
        raise HardsimConfigurationError("idempotency_key must be <= 256 characters")
    return key


def _parse_s3_uri(uri: str) -> tuple[str, str]:
    parsed = urlparse(uri)
    if parsed.scheme != "s3" or not parsed.netloc:
        raise HardsimConfigurationError(f"invalid S3 URI: {uri}")
    key = parsed.path.lstrip("/")
    if not key:
        raise HardsimConfigurationError(f"S3 URI missing key: {uri}")
    return parsed.netloc, key


def _safe_filename(value: str) -> str:
    name = Path(value).name.strip()
    if not name:
        raise HardsimConfigurationError("input file name cannot be empty")
    return name


class HardsimClient:
    def __init__(
        self,
        api_key: str,
        base_url: str = DEFAULT_API_URL,
        timeout_s: float = DEFAULT_TIMEOUT_S,
        max_retries: int = DEFAULT_MAX_RETRIES,
        retry_backoff_s: float = DEFAULT_RETRY_BACKOFF_S,
    ) -> None:
        api_key = api_key.strip()
        if not api_key:
            raise HardsimConfigurationError("api_key is required")
        if timeout_s <= 0:
            raise HardsimConfigurationError("timeout_s must be > 0")
        if max_retries < 0:
            raise HardsimConfigurationError("max_retries must be >= 0")
        if retry_backoff_s < 0:
            raise HardsimConfigurationError("retry_backoff_s must be >= 0")

        self.base_url = base_url.rstrip("/")
        self.timeout_s = float(timeout_s)
        self.max_retries = int(max_retries)
        self.retry_backoff_s = float(retry_backoff_s)
        self.session = _UrlLibSession()
        self.session.headers.update(
            {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            }
        )

    def _create_s3_client(self):
        try:
            import boto3  # type: ignore
        except ImportError as exc:
            raise HardsimConfigurationError(
                "boto3 is required for upload_input(). install with: pip install boto3"
            ) from exc

        kwargs: dict[str, str] = {}
        region = os.getenv("HARDSIM_INPUT_S3_REGION", "").strip()
        endpoint_url = os.getenv("HARDSIM_INPUT_S3_ENDPOINT_URL", "").strip()
        if region:
            kwargs["region_name"] = region
        if endpoint_url:
            kwargs["endpoint_url"] = endpoint_url
        return boto3.client("s3", **kwargs)

    @classmethod
    def from_env(cls) -> "HardsimClient":
        api_key = os.getenv("HARDSIM_API_KEY", "").strip()
        if not api_key:
            raise HardsimConfigurationError("HARDSIM_API_KEY is required")

        base_url = os.getenv("HARDSIM_API_URL", DEFAULT_API_URL).strip() or DEFAULT_API_URL
        timeout_s = float(os.getenv("HARDSIM_HTTP_TIMEOUT_S", str(DEFAULT_TIMEOUT_S)))
        max_retries = int(os.getenv("HARDSIM_HTTP_RETRIES", str(DEFAULT_MAX_RETRIES)))
        retry_backoff_s = float(os.getenv("HARDSIM_HTTP_BACKOFF_S", str(DEFAULT_RETRY_BACKOFF_S)))
        return cls(
            api_key=api_key,
            base_url=base_url,
            timeout_s=timeout_s,
            max_retries=max_retries,
            retry_backoff_s=retry_backoff_s,
        )

    def _request(
        self,
        method: str,
        path_or_url: str,
        *,
        json_body: dict[str, Any] | None = None,
        headers: dict[str, str] | None = None,
        timeout_s: float | None = None,
        retryable: bool | None = None,
        include_auth_header: bool = True,
    ) -> _SimpleResponse:
        method = method.upper()
        url = path_or_url if _is_absolute_url(path_or_url) else f"{self.base_url}{path_or_url}"
        request_timeout = self.timeout_s if timeout_s is None else float(timeout_s)
        is_retryable = method in RETRYABLE_METHODS if retryable is None else retryable

        for attempt in range(self.max_retries + 1):
            try:
                response = self.session.request(
                    method=method,
                    url=url,
                    json=json_body,
                    headers=headers,
                    include_default_headers=include_auth_header,
                    timeout=request_timeout,
                )
            except urllib_error.HTTPError as exc:
                response = _to_response_from_http_error(exc)
            except Exception as exc:
                if is_retryable and attempt < self.max_retries:
                    sleep_s = self.retry_backoff_s * (2**attempt)
                    if sleep_s > 0:
                        time.sleep(sleep_s)
                    continue
                raise HardsimTransportError(str(exc)) from exc

            if response.status_code >= 400:
                if (
                    is_retryable
                    and response.status_code in RETRYABLE_STATUS_CODES
                    and attempt < self.max_retries
                ):
                    sleep_s = self.retry_backoff_s * (2**attempt)
                    if sleep_s > 0:
                        time.sleep(sleep_s)
                    continue
                _raise_http_error(response)

            return response

        raise HardsimTransportError("request retry loop exhausted")

    def run(
        self,
        robot: str,
        scene: str,
        num_envs: int,
        steps: int,
        physics_dt: float = 0.005,
        substeps: int = 2,
        seed: int = 42,
        video: bool = True,
        log_format: str = "zarr",
        idempotency_key: str | None = None,
    ) -> Job:
        payload: dict[str, Any] = {
            "job_type": "rollout",
            "robot": {"urdf_uri": robot},
            "scene": {"scene_id": scene},
            "physics": {"dt": physics_dt, "substeps": substeps},
            "execution": {"num_envs": num_envs, "steps": steps, "seed": seed},
            "outputs": {"video": video, "log_format": log_format},
        }
        key = _normalize_optional_idempotency_key(idempotency_key) or f"auto-{uuid.uuid4().hex}"
        headers = {"Idempotency-Key": key} if key is not None else None
        response = self._request(
            "POST",
            "/v0/jobs",
            json_body=payload,
            headers=headers,
            retryable=True,
        )
        data = response.json()
        return Job(client=self, job_id=data["job_id"])

    def submit(
        self,
        robot: str,
        scene: str,
        num_envs: int,
        steps: int,
        physics_dt: float = 0.005,
        substeps: int = 2,
        seed: int = 42,
        video: bool = True,
        log_format: str = "zarr",
        idempotency_key: str | None = None,
    ) -> Job:
        return self.run(
            robot=robot,
            scene=scene,
            num_envs=num_envs,
            steps=steps,
            physics_dt=physics_dt,
            substeps=substeps,
            seed=seed,
            video=video,
            log_format=log_format,
            idempotency_key=idempotency_key,
        )

    def step(
        self,
        robot: str,
        scene: str,
        num_envs: int,
        steps: int,
        physics_dt: float = 0.005,
        substeps: int = 2,
        seed: int = 42,
        video: bool = True,
        log_format: str = "zarr",
        idempotency_key: str | None = None,
    ) -> Job:
        """Alias for run(...), matching the intended cloud-routed step UX."""
        return self.run(
            robot=robot,
            scene=scene,
            num_envs=num_envs,
            steps=steps,
            physics_dt=physics_dt,
            substeps=substeps,
            seed=seed,
            video=video,
            log_format=log_format,
            idempotency_key=idempotency_key,
        )

    def status(self, job_id: str) -> dict[str, Any]:
        response = self._request("GET", f"/v0/jobs/{job_id}")
        return response.json()

    def wait(
        self,
        job_id: str,
        poll_interval_s: float = 2.0,
        timeout_s: float = 1800.0,
        raise_on_error: bool = True,
    ) -> dict[str, Any]:
        return Job(client=self, job_id=job_id).wait(
            poll_interval_s=poll_interval_s,
            timeout_s=timeout_s,
            raise_on_error=raise_on_error,
        )

    def cancel(self, job_id: str) -> dict[str, Any]:
        response = self._request("POST", f"/v0/jobs/{job_id}/cancel", retryable=True)
        return response.json()

    def get_download_url(self, job_id: str, artifact_name: str) -> str:
        response = self._request(
            "GET",
            f"/v0/jobs/{job_id}/artifacts/{artifact_name}/download-url",
        )
        return response.json()["url"]

    def download_artifact(self, job_id: str, artifact_name: str, timeout_s: float = 120.0) -> bytes:
        url = self.get_download_url(job_id=job_id, artifact_name=artifact_name)
        response = self._request(
            "GET",
            url,
            timeout_s=timeout_s,
            retryable=True,
            include_auth_header=False,
        )
        return response.content

    def download(self, job_id: str, output_dir: str) -> list[str]:
        return Job(client=self, job_id=job_id).download(output_dir=output_dir)

    def upload_input(
        self,
        local_path: str,
        *,
        destination_uri: str | None = None,
        bucket: str | None = None,
        key_prefix: str | None = None,
        object_name: str | None = None,
    ) -> str:
        source_path = Path(local_path).expanduser().resolve()
        if not source_path.is_file():
            raise HardsimConfigurationError(f"local input file not found: {source_path}")

        if destination_uri:
            target_bucket, target_key = _parse_s3_uri(destination_uri)
        else:
            target_bucket = (bucket or os.getenv("HARDSIM_INPUT_S3_BUCKET", "")).strip()
            if not target_bucket:
                raise HardsimConfigurationError(
                    "bucket is required (or set HARDSIM_INPUT_S3_BUCKET) when destination_uri is not provided"
                )
            effective_prefix = (key_prefix or os.getenv("HARDSIM_INPUT_S3_PREFIX", "hardsim/inputs")).strip().strip("/")
            file_name = _safe_filename(object_name or source_path.name)
            unique_name = f"{uuid.uuid4().hex}_{file_name}"
            target_key = "/".join(part for part in (effective_prefix, unique_name) if part)

        content_type = mimetypes.guess_type(source_path.name)[0] or "application/octet-stream"
        s3_client = self._create_s3_client()
        try:
            s3_client.upload_file(
                str(source_path),
                target_bucket,
                target_key,
                ExtraArgs={"ContentType": content_type},
            )
        except Exception as exc:
            raise HardsimTransportError(f"failed to upload input asset to s3://{target_bucket}/{target_key}: {exc}") from exc

        return f"s3://{target_bucket}/{target_key}"


def run(
    robot: str,
    scene: str,
    num_envs: int,
    steps: int,
    physics_dt: float = 0.005,
    substeps: int = 2,
    seed: int = 42,
    video: bool = True,
    log_format: str = "zarr",
    idempotency_key: str | None = None,
) -> Job:
    client = HardsimClient.from_env()
    return client.run(
        robot=robot,
        scene=scene,
        num_envs=num_envs,
        steps=steps,
        physics_dt=physics_dt,
        substeps=substeps,
        seed=seed,
        video=video,
        log_format=log_format,
        idempotency_key=idempotency_key,
    )


def submit(
    robot: str,
    scene: str,
    num_envs: int,
    steps: int,
    physics_dt: float = 0.005,
    substeps: int = 2,
    seed: int = 42,
    video: bool = True,
    log_format: str = "zarr",
    idempotency_key: str | None = None,
) -> Job:
    client = HardsimClient.from_env()
    return client.submit(
        robot=robot,
        scene=scene,
        num_envs=num_envs,
        steps=steps,
        physics_dt=physics_dt,
        substeps=substeps,
        seed=seed,
        video=video,
        log_format=log_format,
        idempotency_key=idempotency_key,
    )


def step(
    robot: str,
    scene: str,
    num_envs: int,
    steps: int,
    physics_dt: float = 0.005,
    substeps: int = 2,
    seed: int = 42,
    video: bool = True,
    log_format: str = "zarr",
    idempotency_key: str | None = None,
) -> Job:
    client = HardsimClient.from_env()
    return client.step(
        robot=robot,
        scene=scene,
        num_envs=num_envs,
        steps=steps,
        physics_dt=physics_dt,
        substeps=substeps,
        seed=seed,
        video=video,
        log_format=log_format,
        idempotency_key=idempotency_key,
    )


def wait(
    job_id: str,
    poll_interval_s: float = 2.0,
    timeout_s: float = 1800.0,
    raise_on_error: bool = True,
) -> dict[str, Any]:
    client = HardsimClient.from_env()
    return client.wait(
        job_id=job_id,
        poll_interval_s=poll_interval_s,
        timeout_s=timeout_s,
        raise_on_error=raise_on_error,
    )


def download(job_id: str, output_dir: str) -> list[str]:
    client = HardsimClient.from_env()
    return client.download(job_id=job_id, output_dir=output_dir)


def upload_input(
    local_path: str,
    *,
    destination_uri: str | None = None,
    bucket: str | None = None,
    key_prefix: str | None = None,
    object_name: str | None = None,
) -> str:
    client = HardsimClient.from_env()
    return client.upload_input(
        local_path=local_path,
        destination_uri=destination_uri,
        bucket=bucket,
        key_prefix=key_prefix,
        object_name=object_name,
    )
