"""
AgentBill Wrapper Classes

Explicit wrapper classes for AI providers that give users direct control
over tracking. Alternative to auto-instrumentation for users who prefer
explicit integration.

Usage:
    from agentbill import AgentBillOpenAI, AgentBillAnthropic
    
    # Instead of: client = OpenAI()
    client = AgentBillOpenAI(api_key="ab_xxx", customer_id="cust-123")
    
    # Use normally - tracking is automatic
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": "Hello!"}]
    )
"""

import time
import hashlib
import json
from typing import Any, Dict, Optional, List, Union
from functools import wraps

from .validation import validate_api_key, validate_customer_id
from .exceptions import AgentBillError, BudgetExceededError, ValidationError


# Model pricing (per 1M tokens)
MODEL_PRICING = {
    # OpenAI
    "gpt-4": {"input": 30.0, "output": 60.0},
    "gpt-4-turbo": {"input": 10.0, "output": 30.0},
    "gpt-4-turbo-preview": {"input": 10.0, "output": 30.0},
    "gpt-4o": {"input": 5.0, "output": 15.0},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "gpt-3.5-turbo": {"input": 0.50, "output": 1.50},
    "o1": {"input": 15.0, "output": 60.0},
    "o1-mini": {"input": 3.0, "output": 12.0},
    "o1-preview": {"input": 15.0, "output": 60.0},
    # Anthropic
    "claude-3-opus-20240229": {"input": 15.0, "output": 75.0},
    "claude-3-sonnet-20240229": {"input": 3.0, "output": 15.0},
    "claude-3-haiku-20240307": {"input": 0.25, "output": 1.25},
    "claude-3-5-sonnet-20240620": {"input": 3.0, "output": 15.0},
    "claude-3-5-sonnet-20241022": {"input": 3.0, "output": 15.0},
    "claude-3-5-haiku-20241022": {"input": 1.0, "output": 5.0},
    "claude-sonnet-4-20250514": {"input": 3.0, "output": 15.0},
    # Cohere
    "command-r-plus": {"input": 2.50, "output": 10.0},
    "command-r": {"input": 0.50, "output": 1.50},
    "command": {"input": 1.0, "output": 2.0},
    # Google
    "gemini-1.5-pro": {"input": 3.50, "output": 10.50},
    "gemini-1.5-flash": {"input": 0.075, "output": 0.30},
    "gemini-1.0-pro": {"input": 0.50, "output": 1.50},
    "gemini-2.0-flash": {"input": 0.10, "output": 0.40},
}


def _calculate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
    """Calculate cost based on model and token counts."""
    pricing = MODEL_PRICING.get(model, {"input": 1.0, "output": 2.0})
    input_cost = (input_tokens / 1_000_000) * pricing["input"]
    output_cost = (output_tokens / 1_000_000) * pricing["output"]
    return input_cost + output_cost


def _hash_prompt(messages: Any) -> str:
    """Generate a hash of the prompt for tracking."""
    try:
        content = json.dumps(messages, sort_keys=True, default=str)
        return hashlib.sha256(content.encode()).hexdigest()[:16]
    except Exception:
        return "unknown"


class BaseAgentBillWrapper:
    """Base class for AgentBill wrapper classes."""
    
    def __init__(
        self,
        api_key: str,
        customer_id: Optional[str] = None,
        daily_budget: Optional[float] = None,
        monthly_budget: Optional[float] = None,
        base_url: Optional[str] = None,
        debug: bool = False,
        **kwargs
    ):
        """
        Initialize AgentBill wrapper.
        
        Args:
            api_key: AgentBill API key (ab_xxx)
            customer_id: Customer identifier for tracking
            daily_budget: Optional daily budget limit (enables Cost Guard)
            monthly_budget: Optional monthly budget limit (enables Cost Guard)
            base_url: AgentBill API base URL
            debug: Enable debug logging
            **kwargs: Additional arguments passed to the underlying client
        """
        validate_api_key(api_key)
        if customer_id:
            validate_customer_id(customer_id)
            
        self._agentbill_api_key = api_key
        self._customer_id = customer_id
        self._daily_budget = daily_budget
        self._monthly_budget = monthly_budget
        self._base_url = base_url or "https://bgwyprqxtdreuutzpbgw.supabase.co/functions/v1"
        self._debug = debug
        self._provider_kwargs = kwargs
        
    def _log(self, message: str):
        """Debug logging."""
        if self._debug:
            print(f"[AgentBill] {message}")
            
    async def _validate_budget(self, model: str) -> Dict[str, Any]:
        """Pre-validate request against Cost Guard."""
        if not self._daily_budget and not self._monthly_budget:
            return {"allowed": True, "reason": "no_budget_configured"}
            
        import aiohttp
        
        payload = {
            "api_key": self._agentbill_api_key,
            "customer_id": self._customer_id,
            "model": model,
            "daily_budget": self._daily_budget,
            "monthly_budget": self._monthly_budget,
        }
        
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    f"{self._base_url}/ai-cost-guard-router",
                    json=payload,
                    headers={"Content-Type": "application/json"},
                    timeout=aiohttp.ClientTimeout(total=5)
                ) as resp:
                    result = await resp.json()
                    self._log(f"Cost Guard validation: {result}")
                    return result
        except Exception as e:
            self._log(f"Cost Guard validation failed (allowing): {e}")
            return {"allowed": True, "reason": "validation_failed_open"}
            
    def _validate_budget_sync(self, model: str) -> Dict[str, Any]:
        """Synchronous budget validation."""
        if not self._daily_budget and not self._monthly_budget:
            return {"allowed": True, "reason": "no_budget_configured"}
            
        import requests
        
        payload = {
            "api_key": self._agentbill_api_key,
            "customer_id": self._customer_id,
            "model": model,
            "daily_budget": self._daily_budget,
            "monthly_budget": self._monthly_budget,
        }
        
        try:
            resp = requests.post(
                f"{self._base_url}/ai-cost-guard-router",
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=5
            )
            result = resp.json()
            self._log(f"Cost Guard validation: {result}")
            return result
        except Exception as e:
            self._log(f"Cost Guard validation failed (allowing): {e}")
            return {"allowed": True, "reason": "validation_failed_open"}
            
    def _track_usage_sync(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
        latency_ms: float,
        prompt_hash: str,
        event_name: Optional[str] = None,
        metadata: Optional[Dict] = None,
    ):
        """Track AI usage synchronously."""
        import requests
        
        cost = _calculate_cost(model, input_tokens, output_tokens)
        
        payload = {
            "api_key": self._agentbill_api_key,
            "customer_id": self._customer_id,
            "model": model,
            "provider": self._provider_name,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
            "cost": cost,
            "latency_ms": latency_ms,
            "prompt_hash": prompt_hash,
            "event_name": event_name or "ai_call",
            "metadata": metadata or {},
        }
        
        try:
            resp = requests.post(
                f"{self._base_url}/track-ai-usage",
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=5
            )
            self._log(f"Usage tracked: {resp.status_code}")
        except Exception as e:
            self._log(f"Failed to track usage: {e}")
            
    async def _track_usage(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
        latency_ms: float,
        prompt_hash: str,
        event_name: Optional[str] = None,
        metadata: Optional[Dict] = None,
    ):
        """Track AI usage asynchronously."""
        import aiohttp
        
        cost = _calculate_cost(model, input_tokens, output_tokens)
        
        payload = {
            "api_key": self._agentbill_api_key,
            "customer_id": self._customer_id,
            "model": model,
            "provider": self._provider_name,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
            "cost": cost,
            "latency_ms": latency_ms,
            "prompt_hash": prompt_hash,
            "event_name": event_name or "ai_call",
            "metadata": metadata or {},
        }
        
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    f"{self._base_url}/track-ai-usage",
                    json=payload,
                    headers={"Content-Type": "application/json"},
                    timeout=aiohttp.ClientTimeout(total=5)
                ) as resp:
                    self._log(f"Usage tracked: {resp.status}")
        except Exception as e:
            self._log(f"Failed to track usage: {e}")
            
    @property
    def _provider_name(self) -> str:
        """Override in subclasses."""
        return "unknown"


class AgentBillOpenAI(BaseAgentBillWrapper):
    """
    OpenAI client wrapper with automatic AgentBill tracking.
    
    Usage:
        from agentbill import AgentBillOpenAI
        
        client = AgentBillOpenAI(
            api_key="ab_xxx",
            customer_id="cust-123",
            daily_budget=10.00,  # Optional Cost Guard
            openai_api_key="sk-xxx"  # Or use OPENAI_API_KEY env var
        )
        
        response = client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": "Hello!"}]
        )
    """
    
    def __init__(
        self,
        api_key: str,
        customer_id: Optional[str] = None,
        daily_budget: Optional[float] = None,
        monthly_budget: Optional[float] = None,
        base_url: Optional[str] = None,
        debug: bool = False,
        openai_api_key: Optional[str] = None,
        **kwargs
    ):
        super().__init__(
            api_key=api_key,
            customer_id=customer_id,
            daily_budget=daily_budget,
            monthly_budget=monthly_budget,
            base_url=base_url,
            debug=debug,
            **kwargs
        )
        
        try:
            from openai import OpenAI
        except ImportError:
            raise ImportError("openai package required. Install with: pip install openai")
            
        # Create underlying OpenAI client
        client_kwargs = {k: v for k, v in kwargs.items() if k != "openai_api_key"}
        if openai_api_key:
            client_kwargs["api_key"] = openai_api_key
        self._client = OpenAI(**client_kwargs)
        
        # Wrap the chat completions
        self.chat = _WrappedChatCompletions(self._client.chat, self)
        
        # Pass through other attributes
        self.models = self._client.models
        self.embeddings = self._client.embeddings
        self.files = self._client.files
        self.images = self._client.images
        self.audio = self._client.audio
        self.moderations = self._client.moderations
        
    @property
    def _provider_name(self) -> str:
        return "openai"


class _WrappedChatCompletions:
    """Wrapped chat completions with tracking."""
    
    def __init__(self, chat, wrapper: BaseAgentBillWrapper):
        self._chat = chat
        self._wrapper = wrapper
        self.completions = _WrappedCompletions(chat.completions, wrapper)


class _WrappedCompletions:
    """Wrapped completions.create with tracking."""
    
    def __init__(self, completions, wrapper: BaseAgentBillWrapper):
        self._completions = completions
        self._wrapper = wrapper
        
    def create(
        self,
        *args,
        agentbill_event_name: Optional[str] = None,
        agentbill_metadata: Optional[Dict] = None,
        **kwargs
    ):
        """Create chat completion with automatic tracking."""
        model = kwargs.get("model", "gpt-4")
        messages = kwargs.get("messages", [])
        prompt_hash = _hash_prompt(messages)
        
        # Pre-validate budget
        validation = self._wrapper._validate_budget_sync(model)
        if not validation.get("allowed", True):
            raise BudgetExceededError(
                reason=validation.get("reason", "budget_exceeded"),
                budget_type=validation.get("budget_type", "unknown"),
                limit=validation.get("limit"),
                current_usage=validation.get("current_usage")
            )
        
        # Make the actual call
        start_time = time.time()
        response = self._completions.create(*args, **kwargs)
        latency_ms = (time.time() - start_time) * 1000
        
        # Extract usage
        input_tokens = getattr(response.usage, "prompt_tokens", 0) if response.usage else 0
        output_tokens = getattr(response.usage, "completion_tokens", 0) if response.usage else 0
        
        # Track usage
        self._wrapper._track_usage_sync(
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            latency_ms=latency_ms,
            prompt_hash=prompt_hash,
            event_name=agentbill_event_name,
            metadata=agentbill_metadata,
        )
        
        return response


class AgentBillAnthropic(BaseAgentBillWrapper):
    """
    Anthropic client wrapper with automatic AgentBill tracking.
    
    Usage:
        from agentbill import AgentBillAnthropic
        
        client = AgentBillAnthropic(
            api_key="ab_xxx",
            customer_id="cust-123",
            anthropic_api_key="sk-ant-xxx"
        )
        
        response = client.messages.create(
            model="claude-3-sonnet-20240229",
            max_tokens=1024,
            messages=[{"role": "user", "content": "Hello!"}]
        )
    """
    
    def __init__(
        self,
        api_key: str,
        customer_id: Optional[str] = None,
        daily_budget: Optional[float] = None,
        monthly_budget: Optional[float] = None,
        base_url: Optional[str] = None,
        debug: bool = False,
        anthropic_api_key: Optional[str] = None,
        **kwargs
    ):
        super().__init__(
            api_key=api_key,
            customer_id=customer_id,
            daily_budget=daily_budget,
            monthly_budget=monthly_budget,
            base_url=base_url,
            debug=debug,
            **kwargs
        )
        
        try:
            from anthropic import Anthropic
        except ImportError:
            raise ImportError("anthropic package required. Install with: pip install anthropic")
            
        client_kwargs = {k: v for k, v in kwargs.items() if k != "anthropic_api_key"}
        if anthropic_api_key:
            client_kwargs["api_key"] = anthropic_api_key
        self._client = Anthropic(**client_kwargs)
        
        # Wrap messages
        self.messages = _WrappedAnthropicMessages(self._client.messages, self)
        
    @property
    def _provider_name(self) -> str:
        return "anthropic"


class _WrappedAnthropicMessages:
    """Wrapped Anthropic messages with tracking."""
    
    def __init__(self, messages, wrapper: BaseAgentBillWrapper):
        self._messages = messages
        self._wrapper = wrapper
        
    def create(
        self,
        *args,
        agentbill_event_name: Optional[str] = None,
        agentbill_metadata: Optional[Dict] = None,
        **kwargs
    ):
        """Create message with automatic tracking."""
        model = kwargs.get("model", "claude-3-sonnet-20240229")
        messages = kwargs.get("messages", [])
        prompt_hash = _hash_prompt(messages)
        
        # Pre-validate budget
        validation = self._wrapper._validate_budget_sync(model)
        if not validation.get("allowed", True):
            raise BudgetExceededError(
                reason=validation.get("reason", "budget_exceeded"),
                budget_type=validation.get("budget_type", "unknown"),
                limit=validation.get("limit"),
                current_usage=validation.get("current_usage")
            )
        
        # Make the actual call
        start_time = time.time()
        response = self._messages.create(*args, **kwargs)
        latency_ms = (time.time() - start_time) * 1000
        
        # Extract usage
        input_tokens = getattr(response.usage, "input_tokens", 0) if response.usage else 0
        output_tokens = getattr(response.usage, "output_tokens", 0) if response.usage else 0
        
        # Track usage
        self._wrapper._track_usage_sync(
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            latency_ms=latency_ms,
            prompt_hash=prompt_hash,
            event_name=agentbill_event_name,
            metadata=agentbill_metadata,
        )
        
        return response


class AgentBillCohere(BaseAgentBillWrapper):
    """
    Cohere client wrapper with automatic AgentBill tracking.
    
    Usage:
        from agentbill import AgentBillCohere
        
        client = AgentBillCohere(
            api_key="ab_xxx",
            customer_id="cust-123",
            cohere_api_key="xxx"
        )
        
        response = client.chat(
            model="command-r-plus",
            message="Hello!"
        )
    """
    
    def __init__(
        self,
        api_key: str,
        customer_id: Optional[str] = None,
        daily_budget: Optional[float] = None,
        monthly_budget: Optional[float] = None,
        base_url: Optional[str] = None,
        debug: bool = False,
        cohere_api_key: Optional[str] = None,
        **kwargs
    ):
        super().__init__(
            api_key=api_key,
            customer_id=customer_id,
            daily_budget=daily_budget,
            monthly_budget=monthly_budget,
            base_url=base_url,
            debug=debug,
            **kwargs
        )
        
        try:
            import cohere
        except ImportError:
            raise ImportError("cohere package required. Install with: pip install cohere")
            
        client_kwargs = {k: v for k, v in kwargs.items() if k != "cohere_api_key"}
        if cohere_api_key:
            client_kwargs["api_key"] = cohere_api_key
        self._client = cohere.Client(**client_kwargs)
        self._original_chat = self._client.chat
        
    def chat(
        self,
        *args,
        agentbill_event_name: Optional[str] = None,
        agentbill_metadata: Optional[Dict] = None,
        **kwargs
    ):
        """Chat with automatic tracking."""
        model = kwargs.get("model", "command-r")
        message = kwargs.get("message", "")
        prompt_hash = _hash_prompt(message)
        
        # Pre-validate budget
        validation = self._validate_budget_sync(model)
        if not validation.get("allowed", True):
            raise BudgetExceededError(
                reason=validation.get("reason", "budget_exceeded"),
                budget_type=validation.get("budget_type", "unknown"),
                limit=validation.get("limit"),
                current_usage=validation.get("current_usage")
            )
        
        # Make the actual call
        start_time = time.time()
        response = self._original_chat(*args, **kwargs)
        latency_ms = (time.time() - start_time) * 1000
        
        # Extract usage from meta
        meta = getattr(response, "meta", None)
        tokens = getattr(meta, "tokens", None) if meta else None
        input_tokens = getattr(tokens, "input_tokens", 0) if tokens else 0
        output_tokens = getattr(tokens, "output_tokens", 0) if tokens else 0
        
        # Track usage
        self._track_usage_sync(
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            latency_ms=latency_ms,
            prompt_hash=prompt_hash,
            event_name=agentbill_event_name,
            metadata=agentbill_metadata,
        )
        
        return response
        
    @property
    def _provider_name(self) -> str:
        return "cohere"


class AgentBillGoogleAI(BaseAgentBillWrapper):
    """
    Google AI (Gemini) client wrapper with automatic AgentBill tracking.
    
    Usage:
        from agentbill import AgentBillGoogleAI
        
        client = AgentBillGoogleAI(
            api_key="ab_xxx",
            customer_id="cust-123",
            google_api_key="xxx"
        )
        
        response = client.generate_content(
            model="gemini-1.5-pro",
            contents="Hello!"
        )
    """
    
    def __init__(
        self,
        api_key: str,
        customer_id: Optional[str] = None,
        daily_budget: Optional[float] = None,
        monthly_budget: Optional[float] = None,
        base_url: Optional[str] = None,
        debug: bool = False,
        google_api_key: Optional[str] = None,
        **kwargs
    ):
        super().__init__(
            api_key=api_key,
            customer_id=customer_id,
            daily_budget=daily_budget,
            monthly_budget=monthly_budget,
            base_url=base_url,
            debug=debug,
            **kwargs
        )
        
        try:
            import google.generativeai as genai
        except ImportError:
            raise ImportError("google-generativeai package required. Install with: pip install google-generativeai")
            
        if google_api_key:
            genai.configure(api_key=google_api_key)
        self._genai = genai
        
    def generate_content(
        self,
        model: str = "gemini-1.5-pro",
        contents: Any = None,
        agentbill_event_name: Optional[str] = None,
        agentbill_metadata: Optional[Dict] = None,
        **kwargs
    ):
        """Generate content with automatic tracking."""
        prompt_hash = _hash_prompt(contents)
        
        # Pre-validate budget
        validation = self._validate_budget_sync(model)
        if not validation.get("allowed", True):
            raise BudgetExceededError(
                reason=validation.get("reason", "budget_exceeded"),
                budget_type=validation.get("budget_type", "unknown"),
                limit=validation.get("limit"),
                current_usage=validation.get("current_usage")
            )
        
        # Make the actual call
        start_time = time.time()
        model_instance = self._genai.GenerativeModel(model)
        response = model_instance.generate_content(contents, **kwargs)
        latency_ms = (time.time() - start_time) * 1000
        
        # Extract usage
        usage = getattr(response, "usage_metadata", None)
        input_tokens = getattr(usage, "prompt_token_count", 0) if usage else 0
        output_tokens = getattr(usage, "candidates_token_count", 0) if usage else 0
        
        # Track usage
        self._track_usage_sync(
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            latency_ms=latency_ms,
            prompt_hash=prompt_hash,
            event_name=agentbill_event_name,
            metadata=agentbill_metadata,
        )
        
        return response
        
    @property
    def _provider_name(self) -> str:
        return "google"


class AgentBillBedrock(BaseAgentBillWrapper):
    """
    AWS Bedrock client wrapper with automatic AgentBill tracking.
    
    Usage:
        from agentbill import AgentBillBedrock
        
        client = AgentBillBedrock(
            api_key="ab_xxx",
            customer_id="cust-123",
            region_name="us-east-1"
        )
        
        response = client.invoke_model(
            modelId="anthropic.claude-3-sonnet-20240229-v1:0",
            body={"prompt": "Hello!"}
        )
    """
    
    def __init__(
        self,
        api_key: str,
        customer_id: Optional[str] = None,
        daily_budget: Optional[float] = None,
        monthly_budget: Optional[float] = None,
        base_url: Optional[str] = None,
        debug: bool = False,
        region_name: str = "us-east-1",
        **kwargs
    ):
        super().__init__(
            api_key=api_key,
            customer_id=customer_id,
            daily_budget=daily_budget,
            monthly_budget=monthly_budget,
            base_url=base_url,
            debug=debug,
            **kwargs
        )
        
        try:
            import boto3
        except ImportError:
            raise ImportError("boto3 package required. Install with: pip install boto3")
            
        client_kwargs = {k: v for k, v in kwargs.items() if k != "region_name"}
        self._client = boto3.client("bedrock-runtime", region_name=region_name, **client_kwargs)
        
    def invoke_model(
        self,
        modelId: str,
        body: Union[str, Dict],
        agentbill_event_name: Optional[str] = None,
        agentbill_metadata: Optional[Dict] = None,
        **kwargs
    ):
        """Invoke model with automatic tracking."""
        prompt_hash = _hash_prompt(body)
        
        # Pre-validate budget
        validation = self._validate_budget_sync(modelId)
        if not validation.get("allowed", True):
            raise BudgetExceededError(
                reason=validation.get("reason", "budget_exceeded"),
                budget_type=validation.get("budget_type", "unknown"),
                limit=validation.get("limit"),
                current_usage=validation.get("current_usage")
            )
        
        # Prepare body
        if isinstance(body, dict):
            import json
            body = json.dumps(body)
        
        # Make the actual call
        start_time = time.time()
        response = self._client.invoke_model(modelId=modelId, body=body, **kwargs)
        latency_ms = (time.time() - start_time) * 1000
        
        # Parse response for usage (varies by model)
        response_body = json.loads(response["body"].read())
        
        # Try to extract usage (Anthropic models on Bedrock)
        usage = response_body.get("usage", {})
        input_tokens = usage.get("input_tokens", 0)
        output_tokens = usage.get("output_tokens", 0)
        
        # Track usage
        self._track_usage_sync(
            model=modelId,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            latency_ms=latency_ms,
            prompt_hash=prompt_hash,
            event_name=agentbill_event_name,
            metadata=agentbill_metadata,
        )
        
        return response_body
        
    @property
    def _provider_name(self) -> str:
        return "bedrock"
