import json
import logging
import os
from pathlib import Path
from typing import Any, AsyncGenerator, Dict, List, Optional, Union
from urllib.parse import urljoin

import httpx
from autobyteus.llm.utils.media_payload_formatter import media_source_to_data_uri

logger = logging.getLogger(__name__)


class CertificateError(Exception):
    """Custom exception for certificate-related errors."""


class AutobyteusClient:
    """Async + sync HTTP client for talking to an Autobyteus LLM server."""

    DEFAULT_SERVER_URL = "https://api.autobyteus.com"
    API_KEY_HEADER = "AUTOBYTEUS_API_KEY"
    API_KEY_ENV_VAR = "AUTOBYTEUS_API_KEY"
    SSL_CERT_FILE_ENV_VAR = "AUTOBYTEUS_SSL_CERT_FILE"

    @staticmethod
    def _get_default_server_url_from_env() -> str:
        hosts = os.getenv("AUTOBYTEUS_LLM_SERVER_HOSTS")
        if not hosts:
            return AutobyteusClient.DEFAULT_SERVER_URL

        first_host = next(
            (host.strip() for host in hosts.split(",") if host.strip()),
            None,
        )
        return first_host or AutobyteusClient.DEFAULT_SERVER_URL

    def __init__(self, server_url: Optional[str] = None):
        """
        Initialize the client.

        Args:
            server_url: Explicit server URL. Takes precedence over env vars.
        """
        self.server_url = server_url or self._get_default_server_url_from_env()
        self.api_key = os.getenv(self.API_KEY_ENV_VAR)

        if not self.api_key:
            raise ValueError(
                f"{self.API_KEY_ENV_VAR} environment variable is required. "
                "Please set it before initializing the client."
            )

        custom_cert_path_str = os.getenv(self.SSL_CERT_FILE_ENV_VAR)
        verify_param: Union[str, bool, Path]

        if custom_cert_path_str:
            custom_cert_path = Path(custom_cert_path_str)
            if not custom_cert_path.exists():
                raise CertificateError(
                    f"Custom SSL certificate file specified via {self.SSL_CERT_FILE_ENV_VAR} "
                    f"not found at: {custom_cert_path}"
                )
            if not custom_cert_path.is_file():
                raise CertificateError(
                    f"Custom SSL certificate path specified via {self.SSL_CERT_FILE_ENV_VAR} "
                    f"is not a file: {custom_cert_path}"
                )
            verify_param = str(custom_cert_path)
            logger.info(
                "Using custom SSL certificate file for TLS verification: %s. "
                "This is the recommended secure method for servers with self-signed or "
                "private CA certificates.",
                verify_param,
            )
        else:
            verify_param = False
            logger.warning(
                "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n"
                "SECURITY WARNING: SSL certificate verification is DISABLED because the \n"
                f"'{self.SSL_CERT_FILE_ENV_VAR}' environment variable is not set.\n"
                "This configuration is INSECURE and makes the client vulnerable to \n"
                "Man-in-the-Middle (MitM) attacks. It should ONLY be used for development \n"
                "or testing in trusted environments with self-signed certificates if \n"
                "providing the certificate path is not possible.\n"
                "FOR PRODUCTION or secure environments with self-signed certificates, \n"
                f"it is STRONGLY RECOMMENDED to set the '{self.SSL_CERT_FILE_ENV_VAR}' \n"
                "environment variable to the path of the server's certificate (.pem file) \n"
                "to enable proper TLS verification.\n"
                "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
            )

        timeout_config = httpx.Timeout(connect=10.0, read=None, write=None, pool=None)

        try:
            self.async_client = httpx.AsyncClient(
                verify=verify_param,
                headers={self.API_KEY_HEADER: self.api_key},
                timeout=timeout_config,
            )
            self.sync_client = httpx.Client(
                verify=verify_param,
                headers={self.API_KEY_HEADER: self.api_key},
                timeout=timeout_config,
            )
        except Exception as exc:
            logger.error(
                "Failed to initialize httpx client with SSL configuration (verify='%s'): %s",
                verify_param,
                exc,
            )
            raise RuntimeError(f"HTTP client initialization failed: {exc}") from exc

        logger.info("Initialized Autobyteus client with server URL: %s", self.server_url)

    @staticmethod
    def _wrap_http_error(exc: httpx.HTTPStatusError) -> RuntimeError:
        """
        Produce a RuntimeError that preserves status code and server-provided details.
        """
        response = exc.response
        detail: str = ""
        # Prefer JSON "detail" if present
        try:
            json_body = response.json()
            detail = json_body.get("detail") or json.dumps(json_body)
        except Exception:
            # Fall back to raw text
            detail = response.text

        message = f"HTTP {response.status_code} {response.reason_phrase}"
        if detail:
            message = f"{message}: {detail}"
        else:
            message = f"{message}: {exc}"
        err = RuntimeError(message)
        err.__cause__ = exc
        return err

    async def get_available_llm_models(self) -> Dict[str, Any]:
        """Async discovery of available LLM models."""
        try:
            response = await self.async_client.get(
                urljoin(self.server_url, "/models/llm")
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Async LLM model fetch error: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Async LLM model fetch error: %s", exc)
            raise RuntimeError(str(exc)) from exc

    def get_available_llm_models_sync(self) -> Dict[str, Any]:
        """Synchronous discovery of available LLM models."""
        try:
            response = self.sync_client.get(urljoin(self.server_url, "/models/llm"))
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Sync LLM model fetch error: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Sync LLM model fetch error: %s", exc)
            raise RuntimeError(str(exc)) from exc

    async def get_available_image_models(self) -> Dict[str, Any]:
        """Async discovery of available image models."""
        try:
            response = await self.async_client.get(
                urljoin(self.server_url, "/models/image")
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Async image model fetch error: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Async image model fetch error: %s", exc)
            raise RuntimeError(str(exc)) from exc

    def get_available_image_models_sync(self) -> Dict[str, Any]:
        """Synchronous discovery of available image models."""
        try:
            response = self.sync_client.get(
                urljoin(self.server_url, "/models/image")
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Sync image model fetch error: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Sync image model fetch error: %s", exc)
            raise RuntimeError(str(exc)) from exc

    async def get_available_audio_models(self) -> Dict[str, Any]:
        """Async discovery of available audio models."""
        try:
            response = await self.async_client.get(
                urljoin(self.server_url, "/models/audio")
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Async audio model fetch error: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Async audio model fetch error: %s", exc)
            raise RuntimeError(str(exc)) from exc

    def get_available_audio_models_sync(self) -> Dict[str, Any]:
        """Synchronous discovery of available audio models."""
        try:
            response = self.sync_client.get(
                urljoin(self.server_url, "/models/audio")
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Sync audio model fetch error: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Sync audio model fetch error: %s", exc)
            raise RuntimeError(str(exc)) from exc

    async def send_message(
        self,
        conversation_id: str,
        model_name: str,
        user_message: str,
        image_urls: Optional[List[str]] = None,
        audio_urls: Optional[List[str]] = None,
        video_urls: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        """Send a message and get a response."""
        try:
            normalized_image_urls = await self._normalize_media_sources(image_urls)
            normalized_audio_urls = await self._normalize_media_sources(audio_urls)
            normalized_video_urls = await self._normalize_media_sources(video_urls)
            data = {
                "conversation_id": conversation_id,
                "model_name": model_name,
                "user_message": user_message,
                "image_urls": normalized_image_urls,
                "audio_urls": normalized_audio_urls,
                "video_urls": normalized_video_urls,
            }
            response = await self.async_client.post(
                urljoin(self.server_url, "/send-message"),
                json=data,
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Error sending message: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Error sending message: %s", exc)
            raise RuntimeError(str(exc)) from exc

    async def stream_message(
        self,
        conversation_id: str,
        model_name: str,
        user_message: str,
        image_urls: Optional[List[str]] = None,
        audio_urls: Optional[List[str]] = None,
        video_urls: Optional[List[str]] = None,
    ) -> AsyncGenerator[Dict[str, Any], None]:
        """Stream a message and get responses."""
        try:
            normalized_image_urls = await self._normalize_media_sources(image_urls)
            normalized_audio_urls = await self._normalize_media_sources(audio_urls)
            normalized_video_urls = await self._normalize_media_sources(video_urls)
            data = {
                "conversation_id": conversation_id,
                "model_name": model_name,
                "user_message": user_message,
                "image_urls": normalized_image_urls,
                "audio_urls": normalized_audio_urls,
                "video_urls": normalized_video_urls,
            }

            async with self.async_client.stream(
                "POST",
                urljoin(self.server_url, "/stream-message"),
                json=data,
            ) as response:
                response.raise_for_status()

                async for line in response.aiter_lines():
                    if line.startswith("data: "):
                        try:
                            chunk = json.loads(line[6:])
                            if "error" in chunk:
                                raise RuntimeError(chunk["error"])
                            yield chunk
                        except json.JSONDecodeError as exc:
                            logger.error("Failed to parse stream chunk: %s", exc)
                            raise RuntimeError("Invalid stream response format") from exc

        except httpx.HTTPStatusError as exc:
            logger.error("Stream error: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Stream error: %s", exc)
            raise RuntimeError(str(exc)) from exc

    async def generate_image(
        self,
        model_name: str,
        prompt: str,
        input_image_urls: Optional[List[str]] = None,
        mask_url: Optional[str] = None,
        generation_config: Optional[Dict[str, Any]] = None,
        session_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Generate or edit an image and return the server response."""
        try:
            normalized_input_image_urls = await self._normalize_media_sources(input_image_urls)
            normalized_mask_url = await self._normalize_single_media_source(mask_url)
            data = {
                "model_name": model_name,
                "prompt": prompt,
                "input_image_urls": normalized_input_image_urls,
                "mask_url": normalized_mask_url,
                "generation_config": generation_config or {},
                "session_id": session_id,
            }
            response = await self.async_client.post(
                urljoin(self.server_url, "/generate-image"),
                json=data,
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Error generating image: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Error generating image: %s", exc)
            raise RuntimeError(str(exc)) from exc

    async def generate_speech(
        self,
        model_name: str,
        prompt: str,
        generation_config: Optional[Dict[str, Any]] = None,
        session_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Generate speech from text and return the server response."""
        try:
            data = {
                "model_name": model_name,
                "prompt": prompt,
                "generation_config": generation_config or {},
                "session_id": session_id,
            }
            response = await self.async_client.post(
                urljoin(self.server_url, "/generate-speech"),
                json=data,
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Error generating speech: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Error generating speech: %s", exc)
            raise RuntimeError(str(exc)) from exc

    async def cleanup(self, conversation_id: str) -> Dict[str, Any]:
        """Clean up a conversation (LLM)."""
        try:
            response = await self.async_client.post(
                urljoin(self.server_url, "/cleanup"),
                json={"conversation_id": conversation_id},
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Cleanup error: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Cleanup error: %s", exc)
            raise RuntimeError(str(exc)) from exc

    async def cleanup_image_session(self, session_id: str) -> Dict[str, Any]:
        """Clean up an image session."""
        try:
            response = await self.async_client.post(
                urljoin(self.server_url, "/cleanup/image"),
                json={"session_id": session_id},
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Image session cleanup error: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Image session cleanup error: %s", exc)
            raise RuntimeError(str(exc)) from exc

    async def cleanup_audio_session(self, session_id: str) -> Dict[str, Any]:
        """Clean up an audio session."""
        try:
            response = await self.async_client.post(
                urljoin(self.server_url, "/cleanup/audio"),
                json={"session_id": session_id},
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as exc:
            logger.error("Audio session cleanup error: %s", exc)
            raise self._wrap_http_error(exc)
        except httpx.HTTPError as exc:
            logger.error("Audio session cleanup error: %s", exc)
            raise RuntimeError(str(exc)) from exc

    async def _normalize_media_sources(
        self, media_sources: Optional[List[str]]
    ) -> List[str]:
        if not media_sources:
            return []

        normalized: List[str] = []
        for source in media_sources:
            if not isinstance(source, str):
                continue
            trimmed = source.strip()
            if not trimmed:
                continue
            normalized.append(await media_source_to_data_uri(trimmed))
        return normalized

    async def _normalize_single_media_source(
        self, media_source: Optional[str]
    ) -> Optional[str]:
        if not isinstance(media_source, str):
            return None
        trimmed = media_source.strip()
        if not trimmed:
            return None
        return await media_source_to_data_uri(trimmed)

    async def close(self) -> None:
        """Close both HTTP clients."""
        await self.async_client.aclose()
        self.sync_client.close()

    async def __aenter__(self) -> "AutobyteusClient":
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:  # type: ignore[override]
        await self.close()
