from __future__ import annotations

from typing import TYPE_CHECKING, Any, Optional, cast

import httpx

if TYPE_CHECKING:
    from .types.generation import Generation

__all__ = [
    "PonyFlashError",
    "APIError",
    "APIBusinessError",
    "APIStatusError",
    "APIConnectionError",
    "APITimeoutError",
    "BadRequestError",
    "AuthenticationError",
    "PermissionDeniedError",
    "NotFoundError",
    "InsufficientCreditsError",
    "ConflictError",
    "FileTooLargeError",
    "RateLimitError",
    "InternalServerError",
    "BadGatewayError",
    "GenerationFailedError",
    "GenerationTimeoutError",
]


class PonyFlashError(Exception):
    pass


class APIBusinessError(PonyFlashError):
    """Raised when the API returns HTTP 200 with a non-zero ``error.code``."""

    code: int
    message: str

    _CODE_LABELS: dict[int, str] = {
        1002: "Business error",
        1006: "Session expired",
    }

    def __init__(self, *, code: int, message: str) -> None:
        label = self._CODE_LABELS.get(code, "")
        prefix = f"[{code}] {label}: " if label else f"[{code}] "
        super().__init__(f"{prefix}{message}")
        self.code = code
        self.message = message


class APIError(PonyFlashError):
    message: str
    request: httpx.Request
    body: object | None
    code: str | None
    request_id: str | None

    def __init__(
        self,
        message: str,
        request: httpx.Request,
        *,
        body: object | None,
    ) -> None:
        super().__init__(message)
        self.message = message
        self.request = request
        self.body = body
        if isinstance(body, dict):
            err = body.get("error", body)
            if isinstance(err, dict):
                self.code = err.get("code")
                self.request_id = err.get("requestId")
            else:
                self.code = None
                self.request_id = None
        else:
            self.code = None
            self.request_id = None


class APIStatusError(APIError):
    response: httpx.Response
    status_code: int

    def __init__(
        self,
        message: str,
        *,
        response: httpx.Response,
        body: object | None,
    ) -> None:
        super().__init__(message, response.request, body=body)
        self.response = response
        self.status_code = response.status_code


class APIConnectionError(APIError):
    def __init__(self, *, message: str = "Connection error.", request: httpx.Request) -> None:
        super().__init__(message, request, body=None)


class APITimeoutError(APIConnectionError):
    def __init__(self, request: httpx.Request) -> None:
        super().__init__(message="Request timed out.", request=request)


class BadRequestError(APIStatusError):
    """400"""


class AuthenticationError(APIStatusError):
    """401"""


class InsufficientCreditsError(APIStatusError):
    """402 — credits not enough. Includes balance info and recharge guidance."""

    balance: int | None
    required: int | None

    def __init__(
        self,
        message: str,
        *,
        response: httpx.Response,
        body: object | None,
    ) -> None:
        self.balance = None
        self.required = None
        if isinstance(body, dict):
            err = body.get("error", body)
            if isinstance(err, dict):
                self.balance = err.get("balance")
                self.required = err.get("required")

        parts = [message]
        if self.balance is not None:
            parts.append(f"Current balance: {self.balance} credits.")
        if self.required is not None:
            parts.append(f"Required: {self.required} credits.")
        parts.append("Recharge via client.account.recharge() to get a payment link.")
        message = " ".join(parts)

        super().__init__(message, response=response, body=body)


class PermissionDeniedError(APIStatusError):
    """403"""


class NotFoundError(APIStatusError):
    """404"""


class RateLimitError(APIStatusError):
    """429"""


class InternalServerError(APIStatusError):
    """500"""


class BadGatewayError(APIStatusError):
    """502"""


class ConflictError(APIStatusError):
    """409 — e.g. FILE_NOT_READY"""


class FileTooLargeError(APIStatusError):
    """413 — FILE_TOO_LARGE"""


class GenerationFailedError(PonyFlashError):
    """The generation task itself failed on the server side."""

    generation: Generation

    def __init__(self, *, generation: Generation) -> None:
        err_msg = ""
        hint = ""
        if generation.error:
            err_msg = f": [{generation.error.code}] {generation.error.message}"
            code = generation.error.code
            if code == "UPSTREAM_MODEL_ERROR":
                hint = " Retry with the same parameters, or try a different model."
            elif code == "INVALID_ARGUMENT":
                hint = " Check the parameters against client.models.get(model_id).supported_modes."
            elif code == "RATE_LIMITED":
                hint = " Wait and retry. Use exponential backoff."
        super().__init__(f"Generation {generation.request_id} failed{err_msg}.{hint}")
        self.generation = generation


class GenerationTimeoutError(PonyFlashError):
    """The SDK gave up polling before the task completed."""

    request_id: str
    timeout: float

    def __init__(self, *, request_id: str, timeout: float) -> None:
        super().__init__(
            f"Generation {request_id} timed out after {timeout}s. "
            f"The task may still be running — call client.generations.get('{request_id}') to check, "
            f"or increase the timeout parameter."
        )
        self.request_id = request_id
        self.timeout = timeout


_STATUS_CODE_MAP: dict[int, type[APIStatusError]] = {
    400: BadRequestError,
    401: AuthenticationError,
    402: InsufficientCreditsError,
    403: PermissionDeniedError,
    404: NotFoundError,
    409: ConflictError,
    413: FileTooLargeError,
    429: RateLimitError,
    500: InternalServerError,
    502: BadGatewayError,
}


def make_status_error(response: httpx.Response, body: object | None) -> APIStatusError:
    message = f"Error code: {response.status_code}"
    if isinstance(body, dict):
        err = body.get("error", body)
        if isinstance(err, dict) and "message" in err:
            message = f"{message} - {err['message']}"

    cls = _STATUS_CODE_MAP.get(response.status_code, APIStatusError)
    return cls(message, response=response, body=body)
