import httpx
from typing import Optional, AsyncIterator, Dict, Any, Iterator, List
import logging
import asyncio
import os
from litserve import LitAPI
from pydantic import BaseModel


class GenerationResponse(BaseModel):
    generated_text: str

class InferenceApi(LitAPI):
    def __init__(self, config: Dict[str, Any]):
        """Initialize the Inference API with configuration."""
        super().__init__()
        self.logger = logging.getLogger(__name__)
        self.logger.info("Initializing Inference API")
        self._device = None
        self.stream = False
        self.config = config
        self.llm_config = config.get('llm_server', {})

    def setup(self, device: Optional[str] = None):
        """Synchronous setup method required by LitAPI"""
        self._device = device
        self.logger.info(f"Inference API setup completed on device: {device}")
        return self  # It's common for setup methods to return self for chaining

    async def _get_client(self):
        """Get or create HTTP client as needed"""
        return httpx.AsyncClient(
            base_url=self.llm_config.get('base_url', 'http://localhost:8001'),
            timeout=float(self.llm_config.get('timeout', 60.0))
        )

    def _get_endpoint(self, endpoint_name: str) -> str:
        """Get full endpoint path including prefix"""
        endpoints = self.llm_config.get('endpoints', {})
        api_prefix = self.llm_config.get('api_prefix', '')
        endpoint = endpoints.get(endpoint_name, '')
        return f"{api_prefix}{endpoint}"

    async def _make_request(
            self,
            method: str,
            endpoint: str,
            *,
            params: Optional[Dict[str, Any]] = None,
            json: Optional[Dict[str, Any]] = None,
            stream: bool = False
    ) -> Any:
        """Make an authenticated request to the LLM Server."""
        access_token = os.environ.get("InferenceAPI")
        headers = {"Authorization": f"Bearer {access_token}"} if access_token else {}
        base_url = self.llm_config.get('base_url', 'http://localhost:8002')
        full_endpoint = f"{base_url.rstrip('/')}/{self._get_endpoint(endpoint).lstrip('/')}"

        try:
            self.logger.info(f"Making {method} request to: {full_endpoint}")
            async with await self._get_client() as client:
                if stream:
                    return await client.stream(
                        method,
                        self._get_endpoint(endpoint),
                        params=params,
                        json=json,
                        headers=headers
                    )
                else:
                    response = await client.request(
                        method,
                        self._get_endpoint(endpoint),
                        params=params,
                        json=json,
                        headers=headers
                    )
                    response.raise_for_status()
                    return response

        except Exception as e:
            self.logger.error(f"Error in request to {full_endpoint}: {str(e)}")
            raise

    def predict(self, x: str, **kwargs) -> Iterator[str]:
        """Non-async prediction method that yields results."""
        loop = asyncio.get_event_loop()
        async def async_gen():
            async for item in self._async_predict(x, **kwargs):
                yield item

        gen = async_gen()
        while True:
            try:
                yield loop.run_until_complete(gen.__anext__())
            except StopAsyncIteration:
                break

    async def _async_predict(self, x: str, **kwargs) -> AsyncIterator[str]:
        """Internal async prediction method."""
        if self.stream:
            async for chunk in self.generate_stream(x, **kwargs):
                yield chunk
        else:
            response = await self.generate_response(x, **kwargs)
            yield response

    async def generate_response(
            self,
            prompt: str,
            system_message: Optional[str] = None,
            max_new_tokens: Optional[int] = None
    ) -> str:
        """Generate a complete response by forwarding the request to the LLM Server."""
        self.logger.debug(f"Forwarding generation request for prompt: {prompt[:50]}...")

        try:
            response = await self._make_request(
                "POST",
                "generate",
                json={
                    "prompt": prompt,
                    "system_message": system_message,
                    "max_new_tokens": max_new_tokens
                }
            )
            data = response.json()
            return data["generated_text"]

        except Exception as e:
            self.logger.error(f"Error in generate_response: {str(e)}")
            raise

    async def generate_stream(
            self,
            prompt: str,
            system_message: Optional[str] = None,
            max_new_tokens: Optional[int] = None
    ) -> AsyncIterator[str]:
        """Generate a streaming response by forwarding the request to the LLM Server."""
        self.logger.debug(f"Forwarding streaming request for prompt: {prompt[:50]}...")

        try:
            async with await self._make_request(
                    "POST",
                    "generate_stream",
                    json={
                        "prompt": prompt,
                        "system_message": system_message,
                        "max_new_tokens": max_new_tokens
                    },
                    stream=True
            ) as response:
                async for chunk in response.aiter_text():
                    yield chunk

        except Exception as e:
            self.logger.error(f"Error in generate_stream: {str(e)}")
            raise

    async def generate_embedding(self, text: str) -> List[float]:
        """Generate embedding vector from input text."""
        self.logger.debug(f"Forwarding embedding request for text: {text[:50]}...")

        try:
            response = await self._make_request(
                "POST",
                "embedding",
                json={"text": text}
            )
            data = response.json()
            return data["embedding"]

        except Exception as e:
            self.logger.error(f"Error in generate_embedding: {str(e)}")
            raise

    async def check_system_status(self) -> Dict[str, Any]:
        """Check system status of the LLM Server."""
        self.logger.debug("Checking system status...")

        try:
            response = await self._make_request(
                "GET",
                "system_status"
            )
            return response.json()

        except Exception as e:
            self.logger.error(f"Error in check_system_status: {str(e)}")
            raise

    async def download_model(self, model_name: Optional[str] = None) -> Dict[str, str]:
        """Download model files from the LLM Server."""
        self.logger.debug(f"Forwarding model download request for: {model_name or 'default model'}")

        try:
            response = await self._make_request(
                "POST",
                "model_download",
                params={"model_name": model_name} if model_name else None
            )
            return response.json()

        except Exception as e:
            self.logger.error(f"Error in download_model: {str(e)}")
            raise

    async def validate_system(self) -> Dict[str, Any]:
        """Validate system configuration and setup."""
        self.logger.debug("Validating system configuration...")

        try:
            response = await self._make_request(
                "GET",
                "system_validate"
            )
            return response.json()

        except Exception as e:
            self.logger.error(f"Error in validate_system: {str(e)}")
            raise

    async def initialize_model(self, model_name: Optional[str] = None) -> Dict[str, Any]:
        """Initialize specified model or default model."""
        self.logger.debug(f"Initializing model: {model_name or 'default'}")

        try:
            response = await self._make_request(
                "POST",
                "model_initialize",
                params={"model_name": model_name} if model_name else None
            )
            return response.json()

        except Exception as e:
            self.logger.error(f"Error in initialize_model: {str(e)}")
            raise

    async def initialize_embedding_model(self, model_name: Optional[str] = None) -> Dict[str, Any]:
        """Initialize embedding model."""
        self.logger.debug(f"Initializing embedding model: {model_name or 'default'}")

        try:
            response = await self._make_request(
                "POST",
                "model_initialize_embedding",
                json={"model_name": model_name} if model_name else {}
            )
            return response.json()

        except Exception as e:
            self.logger.error(f"Error in initialize_embedding_model: {str(e)}")
            raise

    def decode_request(self, request: Any, **kwargs) -> str:
        """Convert the request payload to input format."""
        if isinstance(request, dict) and "prompt" in request:
            return request["prompt"]
        return request

    def encode_response(self, output: Iterator[str], **kwargs) -> Dict[str, Any]:
        """Convert the model output to a response payload."""
        if self.stream:
            return {"generated_text": output}
        try:
            result = next(output)
            return {"generated_text": result}
        except StopIteration:
            return {"generated_text": ""}

    async def cleanup(self):
        """Cleanup method - no longer needed as clients are created per-request"""
        pass