import asyncio
import aiohttp
from typing import Dict, Any, Optional, List
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,
)
from ..exceptions import (
    ProviderError,
    AuthenticationError,
    RateLimitError,
    ModelNotFoundError,
    ConnectionError,
    OperationTimeoutError,
)


class OpenAIProvider:
    """OpenAI provider implementation"""

    def __init__(self, config: ProviderConfig):
        self.config = config
        self.base_url = config.base_url or "https://api.openai.com/v1"
        self.session: Optional[aiohttp.ClientSession] = None
        self.api_key = config.api_key
        self.timeout = config.timeout
        self.max_retries = config.max_retries

    async def initialize(self) -> None:
        """Initialize the provider"""
        if self.session is None:
            print(f"Initializing OpenAI provider, api_key: {'*' * len(self.api_key) if self.api_key else 'None'}")
            self.session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=self.timeout)
            )

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

    async def _make_request(self, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]:
        """Make a request to the OpenAI API"""
        url = f"{self.base_url}/{endpoint}"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        for attempt in range(self.max_retries):
            try:
                async with self.session.post(
                    url,
                    json=data,
                    headers=headers
                ) as response:
                    if response.status == 401:
                        raise AuthenticationError("Invalid API key")
                    if response.status == 429:
                        raise RateLimitError("Rate limit exceeded")
                    if response.status == 404:
                        raise ModelNotFoundError(f"Model not found: {data.get('model')}")
                    if response.status >= 500:
                        raise ProviderError(f"Server error: {response.status}")

                    return await response.json()

            except (aiohttp.ClientConnectionError, asyncio.TimeoutError) as e:
                if attempt == self.max_retries - 1:
                    raise ConnectionError(f"Connection failed: {e}")
                await asyncio.sleep(2 ** attempt)

            except aiohttp.ClientResponseError as e:
                if e.status == 401:
                    raise AuthenticationError("Invalid API key")
                elif e.status == 429:
                    raise RateLimitError("Rate limit exceeded")
                elif e.status == 404:
                    raise ModelNotFoundError(f"Model not found: {data.get('model')}")
                else:
                    raise ProviderError(f"API error: {e}")

    async def make_request(self, request: Request) -> Response:
        """Make a request to the OpenAI API"""
        if request.type.value in ['text', 'text_generation']:
            return await self._generate_text(request)
        elif request.type == ModelType.CHAT:
            return await self._chat(request)
        elif request.type == ModelType.EMBEDDING:
            return await self._generate_embedding(request)
        else:
            raise ValueError(f"Unsupported model type: {request.type}")

    async def _generate_text(self, request: Request) -> Response:
        """Generate text using OpenAI API"""
        data = {
            "model": request.model,
            "prompt": request.prompt,
            "temperature": request.temperature,
            "max_tokens": request.max_tokens,  # Using the field from Request class
            "top_p": 1.0,  # Default value
            "frequency_penalty": 0.0,  # Default value
            "presence_penalty": 0.0,  # Default value
            "stream": False,  # Default value
        }

        if "messages" in request.prompt:
            data["messages"] = request.prompt.split("\n")
            data["messages"] = request.parameters["messages"]

        result = await self._make_request("completions", data)

        return Response(
            request_id=request.id,
            content=result.get("choices", [{}])[0].get("text", ""),
            metadata={
                "model": result.get("model"),
                "usage": result.get("usage", {}),
                "created_at": datetime.utcnow().isoformat()
            }
        )

    async def _chat(self, request: Request) -> Response:
        """Chat using OpenAI API"""
        data = {
            "model": request.model,
            "messages": request.parameters.get("messages", [{
                "role": "user",
                "content": request.prompt
            }]),
            "temperature": request.parameters.get("temperature", 0.7),
            "max_tokens": request.parameters.get("max_tokens"),
            "top_p": request.parameters.get("top_p", 1.0),
            "frequency_penalty": request.parameters.get("frequency_penalty", 0.0),
            "presence_penalty": request.parameters.get("presence_penalty", 0.0),
        }

        result = await self._make_request("chat/completions", data)

        return Response(
            request_id=request.id,
            content=result.get("choices", [{}])[0].get("message", {}).get("content", ""),
            metadata={
                "model": result.get("model"),
                "usage": result.get("usage", {}),
                "created_at": datetime.utcnow().isoformat()
            }
        )

    async def _generate_embedding(self, request: Request) -> Response:
        """Generate embeddings using OpenAI API"""
        data = {
            "model": request.model,
            "input": request.prompt,
        }

        result = await self._make_request("embeddings", data)

        return Response(
            request_id=request.id,
            content=result.get("data", [{}])[0].get("embedding", []),
            metadata={
                "model": result.get("model"),
                "usage": result.get("usage", {}),
                "created_at": datetime.utcnow().isoformat()
            }
        )

    async def get_model_info(self) -> Dict[str, Any]:
        """Get information about available models"""
        try:
            result = await self._make_request("models", {})
            return result
        except Exception:
            return {"error": "Failed to retrieve model information"}

    async def get_usage(self) -> Dict[str, Any]:
        """Get usage information for the API key"""
        try:
            result = await self._make_request("usage", {})
            return result
        except Exception:
            return {"error": "Failed to retrieve usage information"}

    async def list_models(self) -> List[str]:
        """List available models"""
        try:
            result = await self._make_request("models", {})
            return [model["id"] for model in result.get("data", [])]
        except Exception:
            return []

    async def check_health(self) -> bool:
        """Check if the provider is healthy"""
        try:
            await self._make_request("health", {})
            return True
        except Exception:
            return False

    async def get_rate_limit_info(self) -> Dict[str, Any]:
        """Get rate limit information"""
        try:
            # OpenAI doesn't provide direct rate limit info via API
            # This would need to be tracked from response headers
            return {
                "rate_limit": "not_provided_by_api",
                "max_requests_per_minute": "varies_by_tier",
                "max_tokens_per_minute": "varies_by_tier"
            }
        except Exception:
            return {"error": "Failed to retrieve rate limit info"}

    async def get_organization_info(self) -> Dict[str, Any]:
        """Get organization information"""
        try:
            result = await self._make_request("organizations/self", {})
            return result
        except Exception:
            return {"error": "Failed to retrieve organization info"}

    async def create_file(self, file_path: str, purpose: str) -> Dict[str, Any]:
        """Create a file for fine-tuning or other purposes"""
        try:
            with open(file_path, "rb") as f:
                files = {"file": f}
                async with self.session.post(
                    f"{self.base_url}/files",
                    headers={"Authorization": f"Bearer {self.api_key}"},
                    data={"purpose": purpose},
                    timeout=self.timeout
                ) as response:
                    return await response.json()
        except Exception as e:
            raise ProviderError(f"Failed to create file: {e}")

    async def list_files(self) -> List[Dict[str, Any]]:
        """List uploaded files"""
        try:
            result = await self._make_request("files", {})
            return result.get("data", [])
        except Exception:
            return []

    async def delete_file(self, file_id: str) -> bool:
        """Delete an uploaded file"""
        try:
            async with self.session.delete(
                f"{self.base_url}/files/{file_id}",
                headers={"Authorization": f"Bearer {self.api_key}"}
            ) as response:
                return response.status == 200
        except Exception:
            return False