"""
CONXA SDK Client

Main client class for interacting with the CONXA Wallet API.
"""

import time
import json
import warnings
from typing import Optional, Dict, Any
from datetime import datetime

import requests

from .models import (
    QRCodeData,
    SessionStatus,
    ChargeResult,
    WalletBalance,
    ConnectionStatus,
)
from .exceptions import (
    CONXAError,
    AuthenticationError,
    InsufficientBalanceError,
    SessionExpiredError,
    ConnectionTimeoutError,
    SessionNotFoundError,
    RateLimitError,
    APIError,
    ValidationError,
)
from .qr import generate_qr_code, generate_qr_svg


class CONXAClient:
    """
    CONXA SDK Client for AI providers to integrate wallet payments.
    
    This client enables providers to:
    1. Generate QR codes for users to connect their wallet
    2. Poll for user connection status
    3. Charge users for AI usage
    4. Check user balance
    
    Example:
        >>> client = CONXAClient(api_key="pk_live_xxx")
        >>> 
        >>> # Generate QR for user to scan
        >>> qr = client.create_payment_qr(provider_username="user@example.com")
        >>> 
        >>> # Wait for user to connect
        >>> session = client.wait_for_connection(provider_username="user@example.com")
        >>> 
        >>> # Charge for AI usage
        >>> result = client.charge(
        ...     session_token=session.session_token,
        ...     model_name="gpt-4",
        ...     input_tokens=1000,
        ...     output_tokens=500
        ... )
    
    Args:
        api_key: Provider API key (e.g., "pk_live_openai_xxxxx")
        base_url: CONXA API base URL (default: https://api.conxa.in)
        timeout: Request timeout in seconds (default: 30)
        provider_id: Provider ID (optional, extracted from API key if not provided)
        provider_type: Provider type "api" or "web" (default: "api")
    """
    
    DEFAULT_BASE_URL = "https://api.conxa.in"
    DEFAULT_TIMEOUT = 30
    
    def __init__(
        self,
        api_key: str,
        base_url: str = None,
        timeout: int = None,
        provider_id: str = None,
        provider_type: str = "api",
    ):
        if not api_key:
            raise AuthenticationError("API key is required")
        
        self.api_key = api_key
        self.base_url = (base_url or self.DEFAULT_BASE_URL).rstrip("/")
        self.timeout = timeout or self.DEFAULT_TIMEOUT
        self.provider_type = provider_type
        
        # Extract or set provider_id
        # API key format: pk_live_providername_xxxxx
        if provider_id:
            self.provider_id = provider_id
        else:
            # Try to extract from API key (this is a fallback)
            # Providers should ideally set this explicitly
            self.provider_id = self._extract_provider_id_from_key(api_key)
        
        # Session for connection pooling
        self._session = requests.Session()
        self._session.headers.update({
            "Content-Type": "application/json",
            "Accept": "application/json",
            "X-Provider-API-Key": self.api_key,
        })
    
    def _extract_provider_id_from_key(self, api_key: str) -> str:
        """Extract provider ID from API key format"""
        # This is a placeholder - providers should set provider_id explicitly
        # The actual provider_id should come from your registration
        return api_key[:20] if api_key else "unknown"
    
    def _request(
        self,
        method: str,
        endpoint: str,
        data: Dict = None,
        params: Dict = None,
        headers: Dict = None,
    ) -> Dict[str, Any]:
        """
        Make an HTTP request to the CONXA API.
        
        Args:
            method: HTTP method (GET, POST, etc.)
            endpoint: API endpoint (e.g., "/providers/session/status")
            data: Request body data
            params: Query parameters
            headers: Additional headers
        
        Returns:
            Response JSON data
        
        Raises:
            Various CONXA exceptions based on response
        """
        url = f"{self.base_url}{endpoint}"
        
        request_headers = dict(self._session.headers)
        if headers:
            request_headers.update(headers)
        
        try:
            response = self._session.request(
                method=method,
                url=url,
                json=data,
                params=params,
                headers=request_headers,
                timeout=self.timeout,
            )
            
            # Parse response
            try:
                response_data = response.json()
            except json.JSONDecodeError:
                response_data = {"message": response.text}
            
            # Handle errors
            if response.status_code == 401:
                raise AuthenticationError(
                    response_data.get("detail", "Authentication failed")
                )
            elif response.status_code == 402:
                raise InsufficientBalanceError(
                    response_data.get("detail", "Insufficient balance"),
                    current_balance=response_data.get("new_balance"),
                )
            elif response.status_code == 404:
                raise SessionNotFoundError(
                    response_data.get("detail", "Not found")
                )
            elif response.status_code == 422:
                raise ValidationError(
                    response_data.get("message", "Validation error"),
                    errors=response_data.get("errors", []),
                )
            elif response.status_code == 429:
                # Parse Retry-After header (can be seconds as int or HTTP date)
                retry_after = None
                retry_after_header = response.headers.get("Retry-After")
                if retry_after_header:
                    try:
                        # Try parsing as integer (seconds)
                        retry_after = int(retry_after_header)
                    except ValueError:
                        # If not an integer, try parsing as HTTP date
                        try:
                            from email.utils import parsedate_to_datetime
                            retry_date = parsedate_to_datetime(retry_after_header)
                            retry_after = int((retry_date.timestamp() - time.time()))
                            if retry_after < 0:
                                retry_after = None
                        except (ValueError, TypeError):
                            # If parsing fails, use None
                            retry_after = None
                
                raise RateLimitError(
                    response_data.get("detail", "Rate limit exceeded"),
                    retry_after=retry_after,
                )
            elif response.status_code >= 400:
                raise APIError(
                    response_data.get("detail", response_data.get("message", "API error")),
                    status_code=response.status_code,
                    response=response_data,
                )
            
            return response_data
            
        except requests.exceptions.Timeout:
            raise CONXAError("Request timeout")
        except requests.exceptions.ConnectionError:
            raise CONXAError("Connection error - unable to reach CONXA API")
        except requests.exceptions.RequestException as e:
            raise CONXAError(f"Request failed: {str(e)}")
    
    # =========================================================================
    # QR Code Generation
    # =========================================================================
    
    def create_payment_qr(
        self,
        provider_username: str,
        provider_type: str = None,
        limit: Optional[float] = None,
        size: int = 300,
        expires_in: Optional[int] = None,
    ) -> QRCodeData:
        """
        Generate a QR code for user to scan and connect their wallet.
        
        The QR code contains provider information that the CONXA mobile app
        reads to establish a connection.
        
        Args:
            provider_username: User's ID in your system (e.g., email, username)
            provider_type: "api" or "web" (default: instance default)
            limit: Optional spending limit in tokens
            size: QR code image size in pixels (default: 300)
            expires_in: Optional expiration time in seconds. If set, creates a 
                       short-lived session that expires after this duration.
                       Default: None (no expiration)
        
        Returns:
            QRCodeData with:
                - qr_data: JSON string encoded in QR
                - qr_image: PIL Image object (if qrcode installed)
                - qr_base64: Base64 PNG for HTML embedding
                - expires_at: Timestamp when QR expires (if expires_in is set)
        
        Example:
            >>> qr = client.create_payment_qr(
            ...     provider_username="user@example.com",
            ...     limit=10000,
            ...     expires_in=10  # QR expires after 10 seconds
            ... )
            >>> # Display in web app
            >>> html = f'<img src="{qr.qr_base64}" />'
            >>> # Or save to file
            >>> qr.qr_image.save("connect.png")
        """
        qr_data = generate_qr_code(
            provider_id=self.provider_id,
            provider_type=provider_type or self.provider_type,
            provider_username=provider_username,
            limit=limit,
            size=size,
        )
        
        # Track creation time and expiration if expires_in is set
        from datetime import datetime, timezone, timedelta
        qr_data.created_at = datetime.now(timezone.utc)
        
        if expires_in is not None and expires_in > 0:
            qr_data.expires_at = qr_data.created_at + timedelta(seconds=expires_in)
        
        return qr_data
    
    def create_payment_qr_svg(
        self,
        provider_username: str,
        provider_type: str = None,
        limit: Optional[float] = None,
        size: int = 300,
    ) -> str:
        """
        Generate a QR code as SVG string.
        
        Args:
            provider_username: User's ID in your system
            provider_type: "api" or "web"
            limit: Optional spending limit in tokens
            size: SVG size in pixels
        
        Returns:
            SVG string of the QR code
        """
        return generate_qr_svg(
            provider_id=self.provider_id,
            provider_type=provider_type or self.provider_type,
            provider_username=provider_username,
            limit=limit,
            size=size,
        )
    
    # =========================================================================
    # Session / Connection Status
    # =========================================================================
    
    def get_session_status(
        self, 
        provider_username: str,
        retry_on_rate_limit: bool = True,
        max_retries: int = 3,
        qr_data: Optional[QRCodeData] = None,
    ) -> SessionStatus:
        """
        Check if a user has connected their wallet.
        
        Poll this endpoint after showing the QR code to detect when
        the user scans and approves the connection.
        
        Args:
            provider_username: User's ID in your system
            retry_on_rate_limit: If True, automatically retry on rate limit errors with backoff
            max_retries: Maximum number of retries on rate limit (default: 3)
            qr_data: Optional QRCodeData to check if QR has expired
        
        Returns:
            SessionStatus with:
                - status: "pending", "active", "expired", or "not_found"
                - session_token: Token for charges (only if active)
                - expires_at: Session expiration time
        
        Raises:
            RateLimitError: If rate limit exceeded and retries exhausted
            SessionExpiredError: If QR code has expired (when qr_data is provided)
        
        Example:
            >>> qr = client.create_payment_qr("user@example.com", expires_in=10)
            >>> status = client.get_session_status("user@example.com", qr_data=qr)
            >>> if status.is_active:
            ...     print(f"Connected! Token: {status.session_token}")
            >>> elif status.is_pending:
            ...     print("Waiting for user to scan QR...")
        """
        # Check if QR code has expired
        if qr_data and qr_data.is_expired():
            raise SessionExpiredError(
                "QR code has expired. Please generate a new QR code."
            )
        
        last_exception = None
        
        for attempt in range(max_retries + 1):
            try:
                response = self._request(
                    method="GET",
                    endpoint="/providers/session/status",
                    params={"provider_username": provider_username},
                )
                
                return SessionStatus.from_api_response(response)
                
            except RateLimitError as e:
                last_exception = e
                
                if not retry_on_rate_limit or attempt >= max_retries:
                    # Re-raise if we've exhausted retries or retry is disabled
                    raise
                
                # Calculate backoff time
                if e.retry_after:
                    # Use Retry-After header if provided
                    backoff_time = int(e.retry_after)
                else:
                    # Exponential backoff: 2^attempt seconds (2, 4, 8, ...)
                    backoff_time = min(2 ** attempt, 60)  # Cap at 60 seconds
                
                # Wait before retrying
                time.sleep(backoff_time)
        
        # Should never reach here, but just in case
        if last_exception:
            raise last_exception
    
    def wait_for_connection(
        self,
        provider_username: str,
        timeout: int = 120,
        poll_interval: int = 2,
        on_pending: callable = None,
        rate_limit_timeout: int = 10,
        qr_data: Optional[QRCodeData] = None,
    ) -> SessionStatus:
        """
        Wait for user to connect their wallet (blocking).
        
        Polls the session status until the user connects or timeout occurs.
        If rate limited for more than rate_limit_timeout seconds, the QR session
        will be considered expired and user needs to generate a new QR code.
        
        Args:
            provider_username: User's ID in your system
            timeout: Maximum wait time in seconds (default: 120)
            poll_interval: Time between polls in seconds (default: 2)
            on_pending: Optional callback called on each pending poll
            rate_limit_timeout: Seconds of rate limiting before closing QR (default: 10)
            qr_data: Optional QRCodeData to check if QR has expired
        
        Returns:
            SessionStatus with active session
        
        Raises:
            ConnectionTimeoutError: If user doesn't connect within timeout
            SessionExpiredError: If rate limited for more than rate_limit_timeout seconds
        
        Example:
            >>> session = client.wait_for_connection(
            ...     provider_username="user@example.com",
            ...     timeout=60,
            ...     on_pending=lambda s: print("Waiting...")
            ... )
            >>> print(f"Connected! Token: {session.session_token}")
        """
        start_time = time.time()
        rate_limit_start_time = None
        total_rate_limit_time = 0
        
        while True:
            elapsed = time.time() - start_time
            if elapsed >= timeout:
                raise ConnectionTimeoutError(
                    f"User did not connect within {timeout} seconds"
                )
            
            try:
                # Check if QR has expired before polling
                if qr_data and qr_data.is_expired():
                    raise SessionExpiredError(
                        "QR code has expired. Please generate a new QR code."
                    )
                
                status = self.get_session_status(
                    provider_username, 
                    retry_on_rate_limit=True,
                    qr_data=qr_data
                )
                
                # Reset rate limit tracking on successful request
                if rate_limit_start_time is not None:
                    total_rate_limit_time += time.time() - rate_limit_start_time
                    rate_limit_start_time = None
                
                if status.is_active:
                    return status
                
                if status.is_expired:
                    raise SessionExpiredError("Session expired before connection")
                
                if on_pending and status.is_pending:
                    on_pending(status)
                
            except SessionNotFoundError:
                # No session yet, keep waiting
                # Reset rate limit tracking
                if rate_limit_start_time is not None:
                    total_rate_limit_time += time.time() - rate_limit_start_time
                    rate_limit_start_time = None
                pass
            except RateLimitError as e:
                # Track when rate limiting starts
                if rate_limit_start_time is None:
                    rate_limit_start_time = time.time()
                
                # Calculate current total rate limit time
                current_rate_limit_duration = total_rate_limit_time
                if rate_limit_start_time is not None:
                    current_rate_limit_duration += time.time() - rate_limit_start_time
                
                # Check if we've been rate limited for too long
                if current_rate_limit_duration >= rate_limit_timeout:
                    raise SessionExpiredError(
                        f"QR code session expired due to rate limiting. "
                        f"Please generate a new QR code."
                    )
                
                # Handle rate limit by backing off
                if e.retry_after:
                    backoff_time = int(e.retry_after)
                else:
                    # Use exponential backoff, but cap at poll_interval * 4
                    backoff_time = min(poll_interval * 2, 30)
                
                # Log warning but continue waiting
                warnings.warn(
                    f"Rate limit hit while waiting for connection. "
                    f"Backing off for {backoff_time} seconds. "
                    f"Total rate limit time: {current_rate_limit_duration:.1f}s/{rate_limit_timeout}s",
                    UserWarning
                )
                time.sleep(backoff_time)
                continue  # Skip the normal poll_interval sleep
            
            time.sleep(poll_interval)
    
    # =========================================================================
    # Charging API
    # =========================================================================
    
    def charge(
        self,
        session_token: str,
        model_name: str,
        input_tokens: int,
        output_tokens: int,
    ) -> ChargeResult:
        """
        Charge user for AI usage.
        
        Deducts tokens from user's wallet based on model pricing.
        Call this after each AI request to bill the user.
        
        Args:
            session_token: Active session token from connection
            model_name: AI model used (e.g., "gpt-4", "claude-3")
            input_tokens: Number of input tokens used
            output_tokens: Number of output tokens used
        
        Returns:
            ChargeResult with:
                - approved: Whether charge was successful
                - new_balance: User's remaining token balance
        
        Raises:
            InsufficientBalanceError: If user doesn't have enough tokens
            SessionExpiredError: If session has expired
        
        Example:
            >>> result = client.charge(
            ...     session_token="ps_abc123...",
            ...     model_name="gpt-4",
            ...     input_tokens=1500,
            ...     output_tokens=500
            ... )
            >>> if result.approved:
            ...     print(f"Charged! Remaining: {result.new_balance} tokens")
            >>> else:
            ...     print(f"Failed: {result.error}")
        """
        try:
            response = self._request(
                method="POST",
                endpoint="/charge/authorize",
                data={
                    "model_name": model_name,
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens,
                },
                headers={
                    "X-Provider-Session": session_token,
                },
            )
            
            return ChargeResult.from_api_response(response)
            
        except InsufficientBalanceError as e:
            return ChargeResult(
                approved=False,
                new_balance=e.current_balance or 0,
                error=str(e),
            )
        except SessionExpiredError as e:
            return ChargeResult(
                approved=False,
                new_balance=0,
                error=str(e),
            )
    
    def preview_charge(
        self,
        session_token: str,
        model_name: str,
        input_tokens: int,
        output_tokens: int,
    ) -> Dict[str, Any]:
        """
        Preview a charge without deducting tokens.
        
        Use this to check if user can afford a charge before making it.
        
        Args:
            session_token: Active session token
            model_name: AI model to use
            input_tokens: Estimated input tokens
            output_tokens: Estimated output tokens
        
        Returns:
            Dict with estimated cost and affordability
        
        Note:
            This is a client-side estimation. For accurate pricing,
            use the actual charge endpoint.
        """
        # Get current balance
        try:
            # This would need wallet_id from session - simplified here
            return {
                "model_name": model_name,
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "can_proceed": True,  # Simplified - actual implementation needs balance check
            }
        except Exception as e:
            return {
                "can_proceed": False,
                "error": str(e),
            }
    
    # =========================================================================
    # Wallet / Balance
    # =========================================================================
    
    def get_wallet_balance(self, wallet_id: str) -> WalletBalance:
        """
        Get a wallet's token balance.
        
        This is a public endpoint that returns basic wallet info.
        
        Args:
            wallet_id: 16-digit wallet ID
        
        Returns:
            WalletBalance with wallet_id and tokens
        
        Example:
            >>> balance = client.get_wallet_balance("1234567890123456")
            >>> print(f"Balance: {balance.tokens} tokens")
        """
        response = self._request(
            method="GET",
            endpoint=f"/wallet/{wallet_id}",
        )
        
        return WalletBalance.from_api_response(response)
    
    # =========================================================================
    # Utility Methods
    # =========================================================================
    
    def health_check(self) -> bool:
        """
        Check if the CONXA API is reachable.
        
        Returns:
            True if API is healthy, False otherwise
        """
        try:
            response = self._request(
                method="GET",
                endpoint="/health",
            )
            return response.get("status") == "ok"
        except Exception:
            return False
    
    def close(self):
        """Close the HTTP session"""
        self._session.close()
    
    def __enter__(self):
        """Context manager entry"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit"""
        self.close()
        return False
