models/exceptions.py

import time
from typing import Any, List, Optional

from fastapi import HTTPException


class AIStudioProxyError(Exception):
    """Base exception for AIStudio Proxy errors."""
    def __init__(
        self,
        message: str,
        req_id: Optional[str] = None,
        http_status: int = 500,
        retry_after: Optional[int] = None,
        **kwargs: Any
    ):
        self.message = message
        self.req_id = req_id
        self.http_status = http_status
        self.retry_after = retry_after
        self.timestamp = time.time()
        self.context = kwargs

        # Format message with req_id if present for string representation
        super().__init__(self.__str__())

    def __str__(self) -> str:
        if self.req_id:
            return f"[{self.req_id}] {self.message}"
        return self.message

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__} message='{self.message}' req_id='{self.req_id}' http_status={self.http_status} context={self.context}>"

    def to_http_exception(self) -> HTTPException:
        headers = {}
        if self.retry_after is not None:
            headers["Retry-After"] = str(self.retry_after)

        return HTTPException(
            status_code=self.http_status,
            detail=str(self),
            headers=headers if headers else None
        )

# Browser Errors
class BrowserError(AIStudioProxyError):
    def __init__(self, message: str, http_status: int = 503, retry_after: int = 30, **kwargs):
        super().__init__(message, http_status=http_status, retry_after=retry_after, **kwargs)

class BrowserInitError(BrowserError):
    pass

class PageNotReadyError(BrowserError):
    pass

class BrowserCrashedError(BrowserError):
    def __init__(self, message: str = "Browser crashed unexpectedly", **kwargs):
        super().__init__(message, **kwargs)

class SelectorNotFoundError(BrowserError):
    def __init__(self, selector: str, message: str = "Selector not found", **kwargs):
        super().__init__(f"{message}: {selector}", selector=selector, **kwargs)

# Model Errors
class ModelError(AIStudioProxyError):
    def __init__(self, message: str, http_status: int = 422, **kwargs):
        super().__init__(message, http_status=http_status, **kwargs)

class InvalidModelError(ModelError):
    def __init__(self, model_id: str, available_models: Optional[List[str]] = None, message: str = "Invalid model", **kwargs):
        msg = f"{message}: {model_id}"
        if available_models:
            msg += f" (Available: {', '.join(available_models)})"
        super().__init__(msg, model_id=model_id, available_models=available_models, **kwargs)

class ModelSwitchError(ModelError):
    def __init__(self, target_model: str, current_model: str, message: str = "Failed to switch model", **kwargs):
        super().__init__(f"{message} from {current_model} to {target_model}", target_model=target_model, current_model=current_model, **kwargs)

class ModelListError(ModelError):
    pass

# Client Errors
class ClientDisconnectedError(Exception):
    """Client disconnected exception (kept simple for backward compatibility but with extra fields)."""
    def __init__(self, message: str = "Client disconnected", stage: str = "", req_id: Optional[str] = None, http_status: int = 499):
        self.message = message
        self.stage = stage
        self.req_id = req_id
        self.http_status = http_status
        super().__init__(message)

# Validation Errors
class ValidationError(AIStudioProxyError):
    def __init__(self, message: str, http_status: int = 400, **kwargs):
        super().__init__(message, http_status=http_status, **kwargs)

class MissingParameterError(ValidationError):
    def __init__(self, parameter: str, message: str = "Missing parameter", **kwargs):
        super().__init__(f"{message}: {parameter}", parameter=parameter, **kwargs)

class InvalidParameterError(ValidationError):
    def __init__(self, parameter: str, value: Any, reason: str, message: str = "Invalid parameter", **kwargs):
        super().__init__(f"{message} {parameter}={value}: {reason}", parameter=parameter, value=value, reason=reason, **kwargs)

# Stream Errors
class StreamError(AIStudioProxyError):
    def __init__(self, message: str, http_status: int = 502, **kwargs):
        super().__init__(message, http_status=http_status, **kwargs)

class ProxyConnectionError(StreamError):
    def __init__(self, proxy_url: str, message: str = "Failed to connect to proxy", **kwargs):
        super().__init__(f"{message}: {proxy_url}", proxy_url=proxy_url, **kwargs)

class StreamTimeoutError(StreamError):
    def __init__(self, timeout_seconds: float, message: str = "Stream timed out", **kwargs):
        super().__init__(f"{message} after {timeout_seconds}s", timeout_seconds=timeout_seconds, **kwargs)

# Resource Errors
class ResourceError(AIStudioProxyError):
    def __init__(self, message: str, http_status: int = 503, retry_after: int = 60, **kwargs):
        super().__init__(message, http_status=http_status, retry_after=retry_after, **kwargs)

class QueueFullError(ResourceError):
    def __init__(self, queue_size: int, message: str = "Queue full", **kwargs):
        super().__init__(f"{message} (size: {queue_size})", queue_size=queue_size, **kwargs)

# Upstream Errors
class UpstreamError(AIStudioProxyError):
    def __init__(self, message: str, http_status: int = 502, retry_after: int = 10, **kwargs):
        super().__init__(message, http_status=http_status, retry_after=retry_after, **kwargs)

class AIStudioError(UpstreamError):
    def __init__(self, error_message: str, status_code: int, message: str = "AI Studio error", **kwargs):
        super().__init__(f"{message}: {error_message} (Status: {status_code})", ai_studio_status=status_code, error_message=error_message, **kwargs)

class QuotaExceededError(UpstreamError):
    def __init__(self, message: str = "Quota exceeded", retry_after: int = 3600, **kwargs):
        super().__init__(message, retry_after=retry_after, **kwargs)

class EmptyResponseError(UpstreamError):
    def __init__(self, message: str = "Received empty response", **kwargs):
        super().__init__(message, **kwargs)

class QuotaExceededRetry(Exception):
    pass

# Timeout Errors
class TimeoutError(AIStudioProxyError):
    def __init__(self, message: str, http_status: int = 504, **kwargs):
        super().__init__(message, http_status=http_status, **kwargs)

class ResponseTimeoutError(TimeoutError):
    def __init__(self, timeout_seconds: float, message: str = "Response timed out", **kwargs):
        super().__init__(f"{message} after {timeout_seconds}s", timeout_seconds=timeout_seconds, **kwargs)

class ProcessingTimeoutError(TimeoutError):
    def __init__(self, timeout_seconds: Optional[float] = None, message: str = "Processing timeout", **kwargs):
        msg = message
        if timeout_seconds:
            msg += f" after {timeout_seconds}s"
        super().__init__(msg, timeout_seconds=timeout_seconds, **kwargs)

# Configuration Errors
class ConfigurationError(AIStudioProxyError):
    def __init__(self, message: str, http_status: int = 500, **kwargs):
        super().__init__(message, http_status=http_status, **kwargs)

class MissingConfigError(ConfigurationError):
    def __init__(self, config_key: str, message: str = "Missing configuration", **kwargs):
        super().__init__(f"{message}: {config_key}", config_key=config_key, **kwargs)

class InvalidConfigError(ConfigurationError):
    def __init__(self, config_key: str, value: Any, reason: str, message: str = "Invalid configuration", **kwargs):
        super().__init__(f"{message} {config_key}={value}: {reason}", config_key=config_key, value=value, reason=reason, **kwargs)