"""
Anthropic provider implementation for InferenceKit.
Supports Claude models from Anthropic.
"""

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 AnthropicProvider:
    """Anthropic provider implementation for Claude models."""

    def __init__(self, config: ProviderConfig):
        self.config = config
        self.base_url = config.base_url or "https://api.anthropic.com/v1"
        self.session: Optional[aiohttp.ClientSession] = None
        self.api_key = config.api_key
        self.timeout = config.timeout
        self.max_retries = config.max_retries
        self.headers = {
            "x-api-key": self.api_key,
            "anthropic-version": "2023-06-01",
            "Content-Type": "application/json",
        }

    async def initialize(self) -> None:
        """Initialize the provider."""
        if self.session is None:
            print(f"Initializing Anthropic 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 Anthropic API."""
        url = f"{self.base_url}/{endpoint}"

        for attempt in range(self.max_retries):
            try:
                async with self.session.post(
                    url,
                    json=data,
                    headers=self.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 Anthropic 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 Anthropic API."""
        data = {
            "model": request.model,
            "prompt": 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),
            "stream": request.parameters.get("stream", False),
        }

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

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

    async def _chat(self, request: Request) -> Response:
        """Chat using Anthropic 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("messages", data)

        return Response(
            request_id=request.id,
            content=result.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 Anthropic API."""
        # Note: Anthropic doesn't support embeddings directly
        # This would need to use a different approach or model
        raise NotImplementedError("Embeddings are not supported by Anthropic provider")

    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 list_models(self) -> List[str]:
        """List available models."""
        try:
            result = await self._make_request("models", {})
            return [model["id"] for model in result.get("models", [])]
        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:
            # Anthropic provides rate limit info in response headers
            # This would need to be tracked from actual requests
            return {
                "rate_limit": "tracked_from_headers",
                "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("organization", {})
            return result
        except Exception:
            return {"error": "Failed to retrieve organization info"}

    async def create_fine_tuning_job(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Create a fine-tuning job."""
        try:
            result = await self._make_request("fine_tuning/jobs", data)
            return result
        except Exception as e:
            raise ProviderError(f"Failed to create fine-tuning job: {e}")

    async def list_fine_tuning_jobs(self) -> List[Dict[str, Any]]:
        """List fine-tuning jobs."""
        try:
            result = await self._make_request("fine_tuning/jobs", {})
            return result.get("jobs", [])
        except Exception:
            return []

    async def get_fine_tuning_job(self, job_id: str) -> Dict[str, Any]:
        """Get a specific fine-tuning job."""
        try:
            result = await self._make_request(f"fine_tuning/jobs/{job_id}", {})
            return result
        except Exception:
            return {"error": "Failed to get fine-tuning job"}

    async def cancel_fine_tuning_job(self, job_id: str) -> bool:
        """Cancel a fine-tuning job."""
        try:
            async with self.session.post(
                f"{self.base_url}/fine_tuning/jobs/{job_id}/cancel",
                headers=self.headers
            ) as response:
                return response.status == 200
        except Exception:
            return False