"""
Main client implementation for the Paygent SDK.
"""

import json
import logging
import time
from typing import Optional
from urllib.parse import urljoin

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

try:
    import tiktoken
except ImportError:
    tiktoken = None

from .models import (
    UsageData,
    UsageDataWithStrings,
    APIRequest,
    RawUsageData,
    SendUsageV2Response,
    CostBreakdown,
    Customer,
    CustomerCreateOrGetRequest,
)


class Client:
    """Paygent SDK client for tracking usage and costs for AI models."""

    def __init__(self, api_key: str):
        """
        Initialize the Paygent SDK client.

        Args:
            api_key: Your Paygent API key
        """
        self.api_key = api_key
        # Locked configuration - cannot be changed by users
        self.base_url = "https://cp-api.withpaygent.com"
        # self.base_url = "http://localhost:8082"
        self.timeout = 3000
        
        # Setup logging with ERROR level by default (minimal logging)
        self.logger = logging.getLogger(f"paygent_sdk.{id(self)}")
        self.logger.setLevel(logging.ERROR)
        
        # Add console handler if no handlers exist
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)
        
        # Setup HTTP client with retry strategy
        self.session = requests.Session()
        
        # Configure retry strategy
        retry_strategy = Retry(
            total=3,
            backoff_factor=1,
            status_forcelist=[429, 500, 502, 503, 504],
        )
        
        adapter = HTTPAdapter(max_retries=retry_strategy)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)
        
        # Set timeout from locked configuration
        self.session.timeout = self.timeout

    @classmethod
    def new_client(cls, api_key: str) -> 'Client':
        """
        Create a new Paygent SDK client.

        Args:
            api_key: Your Paygent API key
            
        Returns:
            Client instance
        """
        return cls(api_key)


    def send_usage(
        self, 
        agent_id: str, 
        customer_id: str, 
        indicator: str, 
        usage_data: UsageData
    ) -> None:
        """
        Send usage data to the Paygent API.
        
        Args:
            agent_id: Unique identifier for the agent
            customer_id: Unique identifier for the customer
            indicator: Indicator for the usage event
            usage_data: Usage data containing model and token information
            
        Raises:
            requests.RequestException: If the HTTP request fails
            ValueError: If the usage data is invalid
        """
        # Removed verbose logging - only log errors

        # 🎯 AUTOMATIC CACHED TOKEN HANDLING
        # Users can pass total prompt tokens - we automatically subtract cached tokens
        # This makes manual tracking easier (no math required!)
        cached_tokens = usage_data.cached_tokens or 0
        regular_prompt_tokens = usage_data.prompt_tokens - cached_tokens

        # Calculate cost using separated token counts
        adjusted_usage_data = UsageData(
            service_provider=usage_data.service_provider,
            model=usage_data.model,
            prompt_tokens=regular_prompt_tokens,
            completion_tokens=usage_data.completion_tokens,
            total_tokens=usage_data.total_tokens,
            cached_tokens=cached_tokens
        )

        # Calculate cost
        try:
            cost = self._calculate_cost(usage_data.model, adjusted_usage_data)
        except Exception as e:
            self.logger.error(f"Failed to calculate cost: {e}")
            raise ValueError(f"Failed to calculate cost: {e}") from e

        # Cost calculated (no logging for performance)

        # Prepare API request
    def set_log_level(self, level: int) -> None:
        """
        Set the logging level for the client.
        
        Args:
            level: Logging level (e.g., logging.DEBUG, logging.INFO, logging.WARNING, logging.ERROR)
        """
        self.logger.setLevel(level)

    def get_logger(self) -> logging.Logger:
        """
        Get the logger instance for custom logging.
        
        Returns:
            Logger instance
        """
        return self.logger

    def _get_token_count(self, model: str, text: str) -> int:
        """
        Get token count for a given model and text.
        Supports OpenAI, Anthropic, Google, Meta, AWS, Mistral, Cohere, DeepSeek
        
        Args:
            model: The AI model name
            text: Text to count tokens for
            
        Returns:
            Number of tokens
        """
        if not text:
            return 0

        if not tiktoken:
            self.logger.warning("tiktoken not available, using fallback token counting")
            return self._fallback_token_count(text)

        model_lower = model.lower()

        try:
            # OpenAI GPT models
            if model_lower.startswith("gpt-"):
                encoding = tiktoken.encoding_for_model(model)
                return len(encoding.encode(text))
            
            # Anthropic Claude models
            elif model_lower.startswith("claude-"):
                encoding = tiktoken.get_encoding("cl100k_base")
                return len(encoding.encode(text))
            
            # Google DeepMind Gemini models
            elif model_lower.startswith("gemini-"):
                encoding = tiktoken.get_encoding("cl100k_base")
                return len(encoding.encode(text))
            
            # Meta Llama models
            elif model_lower.startswith("llama-"):
                encoding = tiktoken.get_encoding("cl100k_base")
                return len(encoding.encode(text))
            
            # AWS Titan models
            elif model_lower.startswith("titan-"):
                encoding = tiktoken.get_encoding("cl100k_base")
                return len(encoding.encode(text))
            
            # Mistral models
            elif model_lower.startswith("mistral-"):
                encoding = tiktoken.get_encoding("cl100k_base")
                return len(encoding.encode(text))
            
            # Cohere models
            elif model_lower.startswith("command"):
                encoding = tiktoken.get_encoding("cl100k_base")
                return len(encoding.encode(text))
            
            # DeepSeek models
            elif model_lower.startswith("deepseek-"):
                encoding = tiktoken.get_encoding("cl100k_base")
                return len(encoding.encode(text))
            
            # Default fallback
            else:
                self.logger.warning(f"Unknown model '{model}', using cl100k_base encoding")
                encoding = tiktoken.get_encoding("cl100k_base")
                return len(encoding.encode(text))

        except Exception as e:
            self.logger.warning(f"Failed to get token count for model {model}: {e}, using fallback")
            return self._fallback_token_count(text)

    def _fallback_token_count(self, text: str) -> int:
        """
        Fallback token counting method using word-based estimation.
        
        Args:
            text: Text to count tokens for
            
        Returns:
            Estimated number of tokens
        """
        # Simple word-based estimation: 1.3 tokens per word
        word_count = len(text.split())
        return int(word_count * 1.3)

    def send_usage(
        self, 
        agent_id: str, 
        customer_id: str, 
        indicator: str, 
        usage_data: UsageData
    ) -> None:
        """
        Send usage data to the Paygent API.
        This legacy method now internally uses the V2 API for server-side pricing.
        
        Args:
            agent_id: Unique identifier for the agent
            customer_id: Unique identifier for the customer
            indicator: Indicator for the usage event
            usage_data: Usage data containing model and token information
            
        Raises:
            requests.RequestException: If the HTTP request fails
        """
        raw_usage = RawUsageData(
            provider=usage_data.service_provider,
            model=usage_data.model,
            input_tokens=usage_data.prompt_tokens,
            output_tokens=usage_data.completion_tokens,
            cached_tokens=usage_data.cached_tokens,
            plan=usage_data.plan
        )
        self.send_usage_v2(agent_id, customer_id, indicator, raw_usage)

    def send_usage_with_token_string(
        self, 
        agent_id: str, 
        customer_id: str, 
        indicator: str, 
        usage_data: UsageDataWithStrings
    ) -> None:
        """
        Send usage data to the Paygent API using prompt and output strings.
        The function automatically counts tokens using proper tokenizers for each model provider.
        This method now internally uses the V2 API for server-side pricing.
        
        Args:
            agent_id: Unique identifier for the agent
            customer_id: Unique identifier for the customer
            indicator: Indicator for the usage event
            usage_data: Usage data containing prompt and output strings
            
        Raises:
            requests.RequestException: If the HTTP request fails
            ValueError: If the usage data is invalid
        """
        # Calculate token counts for API request
        prompt_tokens = self._get_token_count(usage_data.model, usage_data.prompt_string)
        completion_tokens = self._get_token_count(usage_data.model, usage_data.output_string)

        raw_usage = RawUsageData(
            provider=usage_data.service_provider,
            model=usage_data.model,
            input_tokens=prompt_tokens,
            output_tokens=completion_tokens,
            plan=usage_data.plan
        )
        self.send_usage_v2(agent_id, customer_id, indicator, raw_usage)

    def send_usage_v2(
        self,
        agent_id: str,
        customer_id: str,
        indicator: str,
        raw_usage: RawUsageData
    ) -> SendUsageV2Response:
        """
        Send usage data to the Paygent API v2 (backend-calculated pricing).
        
        Args:
            agent_id: Unique identifier for the agent
            customer_id: Unique identifier for the customer
            indicator: Indicator for the usage event
            raw_usage: Raw usage data including model, provider, tokens, etc.
            
        Returns:
            SendUsageV2Response containing cost breakdown and model metadata
            
        Raises:
            requests.RequestException: If the HTTP request fails
        """
        self.logger.debug(f"Sending V2 usage data for model: {raw_usage.model}")
        
        # Prepare request data
        request_data = {
            "agentId": agent_id,
            "customerId": customer_id,
            "indicator": indicator,
            "rawUsage": {
                "provider": raw_usage.provider,
                "model": raw_usage.model,
                "plan": raw_usage.plan or "default"
            }
        }
        
        # Add optional fields
        if raw_usage.input_tokens is not None:
            request_data["rawUsage"]["inputTokens"] = raw_usage.input_tokens
        if raw_usage.output_tokens is not None:
            request_data["rawUsage"]["outputTokens"] = raw_usage.output_tokens
        if raw_usage.cached_tokens is not None:
            request_data["rawUsage"]["cachedTokens"] = raw_usage.cached_tokens
        if raw_usage.audio_duration is not None:
            request_data["rawUsage"]["audioDuration"] = raw_usage.audio_duration
        if raw_usage.character_count is not None:
            request_data["rawUsage"]["characterCount"] = raw_usage.character_count
        
        # Create HTTP request
        url = urljoin(self.base_url, "/api/v2/usage")
        headers = {
            "Content-Type": "application/json",
            "paygent-api-key": self.api_key
        }
        
        try:
            response = self.session.post(
                url,
                json=request_data,
                headers=headers,
                timeout=30
            )
            
            self.logger.debug(
                f"V2 API response status: {response.status_code}, "
                f"body: {response.text}"
            )
            
            response.raise_for_status()
            
            # Parse response
            response_data = response.json()
            self.logger.info("V2 usage data sent successfully")
            
            # Parse breakdown
            breakdown_data = response_data.get("breakdown", {})
            breakdown = CostBreakdown(
                input_cost=breakdown_data.get("inputCost", 0.0),
                output_cost=breakdown_data.get("outputCost", 0.0),
                cached_cost=breakdown_data.get("cachedCost", 0.0),
                audio_cost=breakdown_data.get("audioCost", 0.0),
                total_cost=breakdown_data.get("totalCost", 0.0)
            )
            
            return SendUsageV2Response(
                cp_data_id=response_data.get("cpDataId", ""),
                calculated_cost=response_data.get("calculatedCost", 0.0),
                breakdown=breakdown,
                model_metadata=response_data.get("modelMetadata", {})
            )
            
        except requests.RequestException as e:
            self.logger.error(f"Failed to send V2 usage data: {e}")
            raise

    def create_or_get_customer(self, request: CustomerCreateOrGetRequest) -> Customer:
        """
        Create a customer if they don't exist, or return the existing customer.
        This method is idempotent - calling it multiple times with the same external_id will return the same customer.
        
        Args:
            request: Customer creation request with required name and external_id
            
        Returns:
            The created or existing customer
            
        Raises:
            ValueError: If validation fails
            requests.RequestException: If the HTTP request fails
        """
        # Validate required fields
        if not request.name or request.name.strip() == '':
            self.logger.error('Validation failed: Customer name is required')
            raise ValueError('Customer name is required')
        
        if not request.external_id or request.external_id.strip() == '':
            self.logger.error('Validation failed: Customer external_id is required')
            raise ValueError('Customer external_id is required')
        
        self.logger.debug(f"Creating or getting customer with external_id: {request.external_id}")
        
        try:
            # Prepare request data
            request_data = {
                "name": request.name,
                "externalId": request.external_id,
            }
            
            # Add optional fields
            if request.email:
                request_data["email"] = request.email
            if request.website:
                request_data["website"] = request.website
            if request.phone:
                request_data["phone"] = request.phone
            if request.address_line1:
                request_data["addressLine1"] = request.address_line1
            if request.address_line2:
                request_data["addressLine2"] = request.address_line2
            if request.city:
                request_data["city"] = request.city
            if request.state:
                request_data["state"] = request.state
            if request.zip_code:
                request_data["zipCode"] = request.zip_code
            if request.country:
                request_data["country"] = request.country
            
            # Call create-or-get endpoint
            url = urljoin(self.base_url, "/api/v1/customers/create-or-get")
            headers = {
                "Content-Type": "application/json",
                "paygent-api-key": self.api_key
            }
            
            response = self.session.post(
                url,
                json=request_data,
                headers=headers,
                timeout=30
            )
            
            response.raise_for_status()
            
            # Parse response
            response_data = response.json()
            self.logger.info(f"Customer created/retrieved successfully with external_id: {request.external_id}")
            
            return Customer(
                id=response_data.get("id", ""),
                external_id=response_data.get("externalId", ""),
                name=response_data.get("name", ""),
                email=response_data.get("email"),
                website=response_data.get("website"),
                phone=response_data.get("phone"),
                address_line1=response_data.get("addressLine1"),
                address_line2=response_data.get("addressLine2"),
                city=response_data.get("city"),
                state=response_data.get("state"),
                zip_code=response_data.get("zipCode"),
                country=response_data.get("country"),
                created_at=response_data.get("createdAt"),
                updated_at=response_data.get("updatedAt")
            )
            
        except requests.RequestException as e:
            self.logger.error(f"Failed to create or get customer: {e}")
            raise
