import asyncio
import aiohttp
from typing import Optional, Dict, Any, List, Union
from datetime import datetime
from uuid import UUID, uuid4
from pydantic import BaseModel, Field
from enum import Enum
from .types import (
    ModelType,
    ProviderType,
    Request,
    Response,
    ProviderConfig,
    ModelConfig,
    BatchRequest,
    BatchResponse,
)
from .exceptions import (
    ModelNotFoundError,
    ProviderError,
    AuthenticationError,
    RateLimitError,
    InvalidRequestError,
    ConnectionError,
    OperationTimeoutError
)
from ..core.cache import SemanticCache
from ..core.streaming import streaming_manager, SyncStreamingWrapper
from ..backends.memory_backend import MemoryBackend


class Model:
    """Main model class for handling AI model inference"""

    def __init__(self, config: ModelConfig):
        self.config = config
        self.provider_type = config.provider_type
        self.model_name = config.model_name
        self.cache = SemanticCache(config.cache_config)
        self.backend = MemoryBackend(config.cache_config)
        self.session: Optional[aiohttp.ClientSession] = None
        self.rate_limit_reset: Optional[datetime] = None
        self.rate_limit_remaining: Optional[int] = None

        # Initialize provider-specific client
        if self.provider_type == ProviderType.OPENAI:
            from inferencekit.providers.openai import OpenAIProvider
            self.provider = OpenAIProvider(config.provider_config)
        else:
            raise ModelNotFoundError(f"Provider {self.provider_type} not supported")

    async def initialize(self) -> None:
        """Initialize the model and establish connections"""
        if self.session is None:
            self.session = aiohttp.ClientSession()
        await self.provider.initialize()

    async def shutdown(self) -> None:
        """Shutdown the model and clean up resources"""
        if self.session:
            await self.session.close()
            self.session = None
        await self.provider.shutdown()

    async def _check_rate_limit(self) -> bool:
        """Check if rate limit has been exceeded"""
        if self.rate_limit_reset and self.rate_limit_remaining == 0:
            now = datetime.utcnow()
            if now < self.rate_limit_reset:
                raise RateLimitError(
                    f"Rate limit exceeded. Reset at: {self.rate_limit_reset}"
                )
        return True

    async def _update_rate_limit(self, headers: Dict[str, str]) -> None:
        """Update rate limit information from response headers"""
        if "x-ratelimit-remaining" in headers:
            self.rate_limit_remaining = int(headers["x-ratelimit-remaining"])
        if "x-ratelimit-reset" in headers:
            reset_timestamp = int(headers["x-ratelimit-reset"])
            self.rate_limit_reset = datetime.utcfromtimestamp(reset_timestamp)

    async def _make_request(self, request: Request) -> Response:
        """Make a request to the provider"""
        print(f"_make_request called with request: {request}")
        print(f"Request type: {request.type}, value: {request.type.value}")
        await self._check_rate_limit()

        try:
            # Check cache first
            similarity_result = self.cache.find_similar(request)
            if similarity_result.best_match and similarity_result.best_score >= 0.9:
                return similarity_result.best_match.response

            # Make actual request
            response = await self.provider.make_request(request)

            # Update rate limit
            await self._update_rate_limit(response.metadata)

            # Cache the response
            self.cache.add_entry(request, response)
            self.backend.save_request(request)
            self.backend.save_response(str(request.id), response)

            return response

        except aiohttp.ClientResponseError as e:
            if e.status == 401:
                raise AuthenticationError("Invalid API credentials")
            elif e.status == 429:
                await self._update_rate_limit(e.headers)
                raise RateLimitError("Rate limit exceeded")
            elif e.status == 404:
                raise ModelNotFoundError(f"Model {self.model_name} not found")
            else:
                raise ProviderError(f"Provider error: {e}")
        except aiohttp.ClientConnectionError:
            raise ConnectionError("Failed to connect to provider")
        except asyncio.TimeoutError:
            raise OperationTimeoutError("Request timed out", self.provider_type.value)
        except Exception as e:
            raise ProviderError(f"Unexpected error: {e}")

    async def generate(self, prompt: str, **kwargs) -> Response:
        """Generate text based on a prompt"""
        request = Request(
            model=self.model_name,
            provider=self.provider_type,
            type=ModelType.TEXT_GENERATION,
            prompt=prompt,
            parameters=kwargs,
        )
        print(f"Created request with type: {request.type}")

        return await self._make_request(request)

    async def chat(self, messages: List[Dict[str, Any]], **kwargs) -> Response:
        """Chat with the model"""
        prompt_lines = []
        for msg in messages:
            role = msg.get("role", "system")
            content = msg.get("content", "")
            prompt_lines.append(f"System: {role}: {content}")

        request = Request(
            model=self.model_name,
            provider=self.provider_type,
            type=ModelType.CHAT,
            prompt="\n".join(prompt_lines),
            parameters=kwargs,
        )

        return await self._make_request(request)

    async def embed(self, text: str, **kwargs) -> Response:
        """Generate embeddings for text"""
        request = Request(
            model=self.model_name,
            provider=self.provider_type,
            type=ModelType.EMBEDDING,
            prompt=text,
            parameters=kwargs,
        )

        return await self._make_request(request)

    async def batch_generate(
        self,
        prompts: List[str],
        batch_size: int = 10,
        max_concurrency: int = 5,
        **kwargs
    ) -> BatchResponse:
        """Generate multiple responses in batches"""
        batch_request = BatchRequest(
            requests=[
                Request(
                    model=self.model_name,
                    provider=self.provider_type,
                    type=ModelType.TEXT_GENERATION,
                    prompt=prompt,
                    parameters=kwargs,
                )
                for prompt in prompts
            ],
            batch_size=batch_size,
            max_concurrency=max_concurrency,
        )

        results = []
        failed_requests = []
        start_time = datetime.utcnow()

        for i in range(0, len(batch_request.requests), batch_request.batch_size):
            batch = batch_request.requests[i:i + batch_request.batch_size]

            tasks = []
            semaphore = asyncio.Semaphore(batch_request.max_concurrency)

            async def bounded_request(request: Request):
                async with semaphore:
                    try:
                        return await self._make_request(request)
                    except Exception:
                        failed_requests.append(request)
                        return None

            for request in batch:
                tasks.append(bounded_request(request))

            batch_results = await asyncio.gather(*tasks)
            results.extend([r for r in batch_results if r is not None])

        total_time = (datetime.utcnow() - start_time).total_seconds()

        return BatchResponse(
            results=results,
            failed_requests=failed_requests,
            total_time=total_time,
        )

    async def get_history(self, limit: int = 100, offset: int = 0) -> list[Dict[str, Any]]:
        """Get request history"""
        return self.backend.list_requests(limit, offset)

    async def search_history(self, query: str, limit: int = 10) -> list[Dict[str, Any]]:
        """Search request history"""
        return self.backend.search_requests(query, limit)

    async def get_cache_stats(self) -> Dict[str, Any]:
        """Get cache statistics"""
        return self.cache.get_stats()

    async def clear_cache(self) -> None:
        """Clear the cache"""
        self.cache.clear()

    async def get_model_info(self) -> Dict[str, Any]:
        """Get information about the model"""
        return await self.provider.get_model_info()

    async def stream_generate(self, prompt: str, **kwargs):
        """Generate text with streaming support."""
        request = Request(
            model=self.model_name,
            provider=self.provider_type,
            type=ModelType.TEXT_GENERATION,
            prompt=prompt,
            parameters=kwargs,
        )

        return streaming_manager.create_stream(request, **kwargs)

    def stream_generate_sync(self, prompt: str, **kwargs):
        """Synchronous wrapper for streaming generation."""
        stream = streaming_manager.stream_generate(prompt, self.model_name, self.provider_type, **kwargs)
        return SyncStreamingWrapper(stream)

    async def stream_chat(self, messages: List[Dict[str, Any]], **kwargs):
        """Chat with streaming support."""
        request = Request(
            model=self.model_name,
            provider=self.provider_type,
            type=ModelType.CHAT,
            prompt="\n".join([
                f"{msg.get('role', 'user')}: {msg.get('content', '')}"
                for msg in messages
            ]),
            parameters=kwargs,
        )

        return streaming_manager.create_stream(request, **kwargs)

    def stream_chat_sync(self, messages: List[Dict[str, Any]], **kwargs):
        """Synchronous wrapper for streaming chat."""
        request = Request(
            model=self.model_name,
            provider=self.provider_type,
            type=ModelType.CHAT,
            prompt="\n".join([
                f"{msg.get('role', 'user')}: {msg.get('content', '')}"
                for msg in messages
            ]),
            parameters=kwargs,
        )

        stream = streaming_manager.create_stream(request, **kwargs)
        return SyncStreamingWrapper(stream)

    async def stream_embeddings(self, text: str, **kwargs):
        """Generate embeddings with streaming support."""
        request = Request(
            model=self.model_name,
            provider=self.provider_type,
            type=ModelType.EMBEDDING,
            prompt=text,
            parameters=kwargs,
        )

        return streaming_manager.create_stream(request, **kwargs)

    def stream_embeddings_sync(self, text: str, **kwargs):
        """Synchronous wrapper for streaming embeddings."""
        request = Request(
            model=self.model_name,
            provider=self.provider_type,
            type=ModelType.EMBEDDING,
            prompt=text,
            parameters=kwargs,
        )

        stream = streaming_manager.create_stream(request, **kwargs)
        return SyncStreamingWrapper(stream)

    async def get_stream_stats(self) -> Dict[str, Any]:
        """Get streaming statistics."""
        return streaming_manager.get_stream_stats()

    def get_active_streams(self) -> List[UUID]:
        """Get list of active stream IDs."""
        return streaming_manager.get_active_streams()

    async def cancel_stream(self, stream_id: UUID) -> bool:
        """Cancel an active stream."""
        return await streaming_manager.cancel_stream(stream_id)

    def get_stream_metrics(self, stream_id: UUID) -> Optional[Dict[str, Any]]:
        """Get metrics for a specific stream."""
        return streaming_manager.get_stream_metrics(stream_id)

    def get_all_metrics(self) -> Dict[str, Any]:
        """Get metrics for all active streams."""
        return streaming_manager.get_all_metrics()

    async def cleanup_streams(self) -> None:
        """Clean up completed or failed streams."""
        await streaming_manager.cleanup_streams()

    async def stream_text_with_progress(
        self,
        prompt: str,
        progress_callback: Optional[Callable[[float], None]] = None,
        **kwargs
    ) -> AsyncGenerator[str, None]:
        """Stream text with progress updates."""
        return await stream_text_with_progress(
            prompt,
            self.model_name,
            self.provider_type,
            progress_callback,
            **kwargs
        )