"""
Patcher for the httpx library (both sync and async).
"""

import json
import logging
import base64
import inspect
from typing import Any, Dict, Optional
import wrapt

from ..rules import _rewrite_url

logger = logging.getLogger(__name__)


class HTTPXPatcher:
    """Patcher for httpx.Client.request and httpx.AsyncClient.request methods."""
    
    _original_sync_request = None
    _original_async_request = None
    _is_patched = False
    _rules: Dict[str, str] = {}
    _verbose: bool = False
    _agent_token: Optional[str] = None
    
    @classmethod
    def patch(cls, rules: Dict[str, str], agent_token: Optional[str] = None, verbose: bool = False) -> None:
        """
        Apply the patch to httpx library.
        
        Args:
            rules: Dictionary mapping hostnames to destinations
            verbose: Whether to log intercepted requests
        """
        if cls._is_patched:
            logger.warning("HTTPX patcher is already active")
            return
        
        try:
            import httpx
            
            # Patch sync client
            cls._original_sync_request = httpx.Client.request
            cls._rules = rules
            cls._verbose = verbose
            cls._agent_token = agent_token
            
            def sync_wrapper(wrapped, instance, args, kwargs):
                return cls._intercept_sync_request(wrapped, instance, args, kwargs)
            
            wrapt.wrap_function_wrapper('httpx', 'Client.request', sync_wrapper)
            
            # Patch async client
            cls._original_async_request = httpx.AsyncClient.request
            
            async def async_wrapper(wrapped, instance, args, kwargs):
                return await cls._intercept_async_request(wrapped, instance, args, kwargs)
            
            wrapt.wrap_function_wrapper('httpx', 'AsyncClient.request', async_wrapper)
            
            cls._is_patched = True
            
            if verbose:
                logger.info("HTTPX library patched successfully (sync and async)")
                
        except ImportError:
            logger.warning("httpx library not available, skipping patch")
        except Exception as e:
            logger.error(f"Failed to patch httpx library: {e}")
            raise
    
    @classmethod
    def unpatch(cls) -> None:
        """Remove the patch from httpx library."""
        if not cls._is_patched:
            return
        
        try:
            import httpx
            if cls._original_sync_request:
                httpx.Client.request = cls._original_sync_request
            if cls._original_async_request:
                httpx.AsyncClient.request = cls._original_async_request
            
            cls._is_patched = False
            cls._original_sync_request = None
            cls._original_async_request = None
            cls._rules = {}
            cls._agent_token = None
            
            if cls._verbose:
                logger.info("HTTPX library unpatched")
        except Exception as e:
            logger.error(f"Failed to unpatch httpx library: {e}")
    
    @classmethod
    def _intercept_sync_request(
        cls,
        wrapped,
        instance,
        args,
        kwargs
    ) -> Any:
        """
        Intercept and rewrite the request URL for sync httpx.Client.
        
        Args:
            wrapped: The original request method
            instance: The Client instance
            args: Positional arguments
            kwargs: Keyword arguments
            
        Returns:
            Response from the original request method
        """
        # Extract URL from args or kwargs
        # httpx.Client.request(method, url, ...) - method is args[0], url is args[1] or kwargs['url']
        url = None
        if 'url' in kwargs:
            url = kwargs['url']
        elif len(args) > 1:
            url = args[1]  # args[0] is method, args[1] is URL
        
        if url:
            original_url = str(url)
            rewritten_url = _rewrite_url(original_url, cls._rules)
            
            if rewritten_url != original_url:
                if cls._verbose:
                    logger.info(
                        f"[Kindred AI] Rewriting request: {original_url} -> {rewritten_url}"
                    )
                
                # Capture original payload/body
                original_payload = None
                if 'data' in kwargs:
                    original_payload = kwargs['data']
                elif 'json' in kwargs:
                    original_payload = kwargs['json']
                elif 'content' in kwargs:
                    original_payload = kwargs['content']
                elif 'files' in kwargs:
                    # For file uploads, we can't easily serialize, so skip
                    original_payload = None
                
                # Initialize headers if not present
                if 'headers' not in kwargs or kwargs['headers'] is None:
                    kwargs['headers'] = {}
                # Ensure headers is a dict
                if not isinstance(kwargs['headers'], dict):
                    # Convert list of tuples to dict if needed
                    if isinstance(kwargs['headers'], (list, tuple)):
                        kwargs['headers'] = dict(kwargs['headers'])
                    else:
                        # For httpx.Headers or other types, try to convert
                        try:
                            kwargs['headers'] = dict(kwargs['headers'])
                        except (TypeError, ValueError):
                            kwargs['headers'] = {}
                
                # Add original destination URL header
                kwargs['headers']['X-Kindred-Original-URL'] = original_url
                
                # Add original payload header if available
                if original_payload is not None:
                    try:
                        # Try to serialize as JSON if it's a dict/list
                        if isinstance(original_payload, (dict, list)):
                            payload_str = json.dumps(original_payload)
                            kwargs['headers']['X-Kindred-Original-Payload'] = payload_str
                            kwargs['headers']['X-Kindred-Original-Payload-Type'] = 'json'
                        elif isinstance(original_payload, str):
                            kwargs['headers']['X-Kindred-Original-Payload'] = original_payload
                            kwargs['headers']['X-Kindred-Original-Payload-Type'] = 'string'
                        elif isinstance(original_payload, bytes):
                            # Encode binary data as base64
                            payload_str = base64.b64encode(original_payload).decode('utf-8')
                            kwargs['headers']['X-Kindred-Original-Payload'] = payload_str
                            kwargs['headers']['X-Kindred-Original-Payload-Type'] = 'base64'
                        else:
                            # For other types, convert to string
                            kwargs['headers']['X-Kindred-Original-Payload'] = str(original_payload)
                            kwargs['headers']['X-Kindred-Original-Payload-Type'] = 'string'
                    except Exception as e:
                        if cls._verbose:
                            logger.warning(f"Failed to serialize original payload: {e}")
                
                # Add agent token header if available
                if cls._agent_token:
                    kwargs['headers']['X-Kindred-Agent-Token'] = cls._agent_token
                    if cls._verbose:
                        logger.debug(f"Added X-Kindred-Agent-Token header to intercepted request")
                
                # Update the URL in args or kwargs
                if 'url' in kwargs:
                    kwargs['url'] = rewritten_url
                elif len(args) > 1:
                    # Preserve method (args[0]) and update URL (args[1])
                    args = (args[0], rewritten_url) + args[2:]
        
        # Call the original method
        return wrapped(*args, **kwargs)
    
    @classmethod
    async def _intercept_async_request(
        cls,
        wrapped,
        instance,
        args,
        kwargs
    ) -> Any:
        """
        Intercept and rewrite the request URL for async httpx.AsyncClient.
        
        Args:
            wrapped: The original request method
            instance: The AsyncClient instance
            args: Positional arguments
            kwargs: Keyword arguments
            
        Returns:
            Response from the original request method
        """
        # Extract URL from args or kwargs
        # httpx.Client.request(method, url, ...) - method is args[0], url is args[1] or kwargs['url']
        url = None
        if 'url' in kwargs:
            url = kwargs['url']
        elif len(args) > 1:
            url = args[1]  # args[0] is method, args[1] is URL
        
        if url:
            original_url = str(url)
            rewritten_url = _rewrite_url(original_url, cls._rules)
            
            if rewritten_url != original_url:
                if cls._verbose:
                    logger.info(
                        f"[Kindred AI] Rewriting request: {original_url} -> {rewritten_url}"
                    )
                
                # Update the URL in args or kwargs
                if 'url' in kwargs:
                    kwargs['url'] = rewritten_url
                elif len(args) > 1:
                    # Preserve method (args[0]) and update URL (args[1])
                    args = (args[0], rewritten_url) + args[2:]
        
        # Call the original method (await if it's async)
        result = wrapped(*args, **kwargs)
        if inspect.isawaitable(result):
            return await result
        return result

