"""AgentBill SDK Client"""
import time
import httpx
from typing import Any, Optional, Dict
from .tracer import AgentBillTracer
from .types import AgentBillConfig
from .pricing import calculate_cost, calculate_image_cost, calculate_audio_cost, is_embedding_model


class AgentBill:
    """
    AgentBill SDK for Python
    
    Example:
        >>> from agentbill import AgentBill
        >>> import openai
        >>> 
        >>> # Initialize AgentBill
        >>> agentbill = AgentBill.init({
        ...     "api_key": "your-api-key",
        ...     "customer_id": "customer-123",
        ...     "debug": True
        ... })
        >>> 
        >>> # Wrap your OpenAI client
        >>> client = agentbill.wrap_openai(openai.OpenAI(api_key="sk-..."))
        >>> 
        >>> # Use normally - all calls are tracked!
        >>> response = client.chat.completions.create(
        ...     model="gpt-4",
        ...     messages=[{"role": "user", "content": "Hello!"}]
        ... )
    """
    
    def __init__(self, config: AgentBillConfig):
        self.config = config
        self.tracer = AgentBillTracer(config)
    
    @classmethod
    def init(cls, config: AgentBillConfig) -> "AgentBill":
        """Initialize AgentBill SDK"""
        return cls(config)
    
    def _estimate_tokens(self, text: str) -> int:
        """Simple token estimation: ~4 chars per token"""
        return max(1, len(str(text)) // 4)
    
    def _estimate_cost(self, model: str, input_tokens: int, output_tokens: int, provider: str = "openai") -> float:
        """Estimate cost using centralized pricing database"""
        return calculate_cost(model, input_tokens, output_tokens, provider)
    
    def _validate_request(self, model: str, messages: Any, estimated_output_tokens: int = 1000) -> Dict:
        """
        Call ai-cost-guard-router edge function (tier-based routing)
        
        CRITICAL: Always validates against company/customer/agent budgets in DB.
        SDK budgets (daily_budget, monthly_budget) are OPTIONAL overrides that add
        stricter limits on top of DB policies - they are NOT required for validation.
        """
        # ALWAYS validate when customer_id is present - backend checks DB policies
        if not self.config.get("customer_id"):
            if self.config.get("debug"):
                print("[AgentBill Cost Guard] ⚠️ No customer_id - skipping validation (not recommended)")
            return {"allowed": True}
        
        url = f"{self.config.get('base_url', 'https://bgwyprqxtdreuutzpbgw.supabase.co')}/functions/v1/ai-cost-guard-router"
        
        payload = {
            "api_key": self.config["api_key"],
            "customer_id": self.config.get("customer_id"),
            "model": model,
            "messages": messages,
            # SDK budgets are OPTIONAL overrides - backend uses DB budgets if not set
            "daily_budget_override": self.config.get("daily_budget"),
            "monthly_budget_override": self.config.get("monthly_budget"),
        }
        
        try:
            with httpx.Client(timeout=10) as client:
                response = client.post(url, json=payload)
                
                # Handle HTTP errors explicitly - don't fail open on 4xx/5xx
                if response.status_code >= 500:
                    error_text = response.text
                    if self.config.get("debug"):
                        print(f"[AgentBill Cost Guard] ❌ Router error {response.status_code}: {error_text}")
                    # Only fail open on server errors (temporary issues)
                    return {"allowed": True, "reason": "Router server error (failed open)"}
                elif response.status_code >= 400:
                    error_text = response.text
                    if self.config.get("debug"):
                        print(f"[AgentBill Cost Guard] ❌ Router rejected {response.status_code}: {error_text}")
                    # Fail closed on client errors (bad request, auth issues)
                    return {"allowed": False, "reason": f"Router error: {error_text}"}
                
                result = response.json()
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] Router response: {result}")
                    if result.get("tier"):
                        print(f"[AgentBill] Tier: {result['tier']}, Mode: {result.get('mode', 'unknown')}")
                    if not result.get("allowed"):
                        print(f"[AgentBill Cost Guard] ❌ BLOCKED: {result.get('reason', 'Unknown reason')}")
                
                return result
        except Exception as e:
            if self.config.get("debug"):
                print(f"[AgentBill Cost Guard] ⚠️ Router network error: {e} (failing open)")
            # Only fail open on network errors (temporary issues)
            return {"allowed": True, "reason": "Router network error (failed open)"}
    
    def _track_usage(self, model: str, provider: str, input_tokens: int, output_tokens: int, latency_ms: float, cost: float, event_name: str = "ai_request", trace_id: str = None, span_id: str = None, request_id: str = None, prompt_metadata: Dict = None):
        """Call track-ai-usage edge function with optional trace correlation and Cost Guard linking"""
        url = f"{self.config.get('base_url', 'https://bgwyprqxtdreuutzpbgw.supabase.co')}/functions/v1/track-ai-usage"
        
        payload = {
            "api_key": self.config["api_key"],
            "customer_id": self.config.get("customer_id"),
            "agent_id": self.config.get("agent_id"),
            "event_name": event_name,
            "model": model,
            "provider": provider,
            "prompt_tokens": input_tokens,
            "completion_tokens": output_tokens,
            "latency_ms": latency_ms,
            "cost": cost,
        }
        
        # Add OTEL correlation if available
        if trace_id:
            payload["trace_id"] = trace_id
        if span_id:
            payload["span_id"] = span_id
        
        # CRITICAL: Add request_id from router validation for audit log linking
        if request_id:
            payload["request_id"] = request_id
        
        # CRITICAL: Add all prompt metadata for AI Economics / Prompt Profitability
        if prompt_metadata:
            if "prompt_hash" in prompt_metadata:
                payload["prompt_hash"] = prompt_metadata["prompt_hash"]
            if "prompt_name" in prompt_metadata:
                payload["prompt_name"] = prompt_metadata["prompt_name"]
            if "prompt_version" in prompt_metadata:
                payload["prompt_version"] = prompt_metadata["prompt_version"]
            if "prompt_owner" in prompt_metadata:
                payload["prompt_owner"] = prompt_metadata["prompt_owner"]
            if "prompt_sample" in prompt_metadata:
                payload["prompt_sample"] = prompt_metadata["prompt_sample"]
        
        # Add OTEL correlation if available
        if trace_id:
            payload["trace_id"] = trace_id
        if span_id:
            payload["span_id"] = span_id
        
        try:
            with httpx.Client(timeout=10) as client:
                client.post(url, json=payload)
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] Usage tracked: ${cost:.4f}")
        except Exception as e:
            if self.config.get("debug"):
                print(f"[AgentBill Cost Guard] Tracking failed: {e}")
    
    def wrap_openai(self, client: Any) -> Any:
        """Wrap OpenAI client with Cost Guard protection"""
        
        # Track chat completions
        original_chat_create = client.chat.completions.create
        def tracked_chat_create(*args, **kwargs):
            model = kwargs.get("model", "unknown")
            messages = kwargs.get("messages", [])
            temperature = kwargs.get("temperature")
            max_tokens = kwargs.get("max_tokens", kwargs.get("max_completion_tokens"))
            
            # Extract event_name from agentbill_options if provided (don't pass to OpenAI)
            agentbill_options = kwargs.pop("agentbill_options", {})
            event_name = agentbill_options.get("event_name", "ai_request")
            
            # Extract prompt metadata for AI Economics analysis
            prompt_metadata = {
                "prompt_hash": agentbill_options.get("prompt_hash"),
                "prompt_name": agentbill_options.get("prompt_name"),
                "prompt_version": agentbill_options.get("prompt_version"),
                "prompt_owner": agentbill_options.get("prompt_owner"),
                "prompt_sample": agentbill_options.get("prompt_sample"),
            }
            # Remove None values
            prompt_metadata = {k: v for k, v in prompt_metadata.items() if v is not None}
            
            if self.config.get("debug"):
                print(f"[AgentBill] Intercepting OpenAI chat.completions.create with model: {model}")
            
            start_time = time.time()
            span = self.tracer.start_span("openai.chat.completions.create", {
                # OTEL GenAI compliant attributes
                "gen_ai.system": "openai",
                "gen_ai.request.model": model,
                "gen_ai.operation.name": "chat",
                # Backward compatibility
                "model": model,
                "provider": "openai"
            })
            
            # CRITICAL: ALWAYS validate before AI calls (regardless of SDK budget config)
            # Backend checks company/customer/agent budgets from DB
            # SDK budgets are optional overrides that ADD stricter limits
            validation_request_id = None
            validation = self._validate_request(model, messages)
            
            # Extract request_id for audit log correlation
            validation_request_id = validation.get("request_id")
            
            # CRITICAL: Block request if not allowed
            if not validation.get("allowed"):
                error_msg = validation.get("reason", "Request blocked by Cost Guard")
                error = Exception(error_msg)
                error.code = "BUDGET_EXCEEDED" if "budget" in error_msg.lower() or "BUDGET" in error_msg else \
                             "RATE_LIMIT_EXCEEDED" if "rate" in error_msg.lower() or "RATE" in error_msg else \
                             "POLICY_VIOLATION"
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ❌ Request BLOCKED: {error_msg}")
                span.set_status(1, str(error))
                span.end()
                raise error
            
            # Validation passed - proceed with OpenAI call
            if self.config.get("debug"):
                print(f"[AgentBill Cost Guard] ✓ Validation passed: {validation.get('mode', 'validated')}")
                if validation_request_id:
                    print(f"[AgentBill Cost Guard] request_id: {validation_request_id}")
            
            # Step 2: Execute OpenAI call (after validation passed or no customer_id set)
            try:
                response = original_chat_create(*args, **kwargs)
                latency = (time.time() - start_time) * 1000
                
                # Step 3: Track usage after successful AI call (with OTEL + Cost Guard correlation)
                input_tokens = response.usage.prompt_tokens
                output_tokens = response.usage.completion_tokens
                cost = self._estimate_cost(model, input_tokens, output_tokens)
                
                if self.config.get("customer_id"):
                    # Pass validation_request_id and prompt_metadata for full traceability
                    self._track_usage(
                        model, "openai", input_tokens, output_tokens, latency, cost, 
                        event_name, span.trace_id, span.span_id, 
                        validation_request_id, prompt_metadata
                    )
                
                span.set_attributes({
                    "gen_ai.usage.input_tokens": input_tokens,
                    "gen_ai.usage.output_tokens": output_tokens,
                    "gen_ai.usage.total_tokens": response.usage.total_tokens,
                    "latency_ms": latency
                })
                span.set_status(0)
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ✓ Protected call completed: ${cost:.4f}")
                
                return response
            except Exception as e:
                span.set_status(1, str(e))
                raise
            finally:
                span.end()
        
        client.chat.completions.create = tracked_chat_create
        
        # Track embeddings
        original_embeddings_create = client.embeddings.create
        def tracked_embeddings_create(*args, **kwargs):
            model = kwargs.get("model", "text-embedding-ada-002")
            
            # Extract event_name from agentbill_options if provided
            agentbill_options = kwargs.pop("agentbill_options", {})
            event_name = agentbill_options.get("event_name", "embedding_request")
            
            if self.config.get("debug"):
                print(f"[AgentBill] Intercepting OpenAI embeddings.create with model: {model}")
            
            start_time = time.time()
            span = self.tracer.start_span("openai.embeddings.create", {
                # OTEL GenAI compliant attributes
                "gen_ai.system": "openai",
                "gen_ai.request.model": model,
                "gen_ai.operation.name": "text_embedding",
                # Backward compatibility
                "model": model,
                "provider": "openai"
            })
            
            try:
                response = original_embeddings_create(*args, **kwargs)
                latency = (time.time() - start_time) * 1000
                
                # ✅ FIX: Track cost for embeddings (was missing!)
                input_tokens = response.usage.prompt_tokens
                cost = self._estimate_cost(model, input_tokens, 0, "openai")
                self._track_usage(model, "openai", input_tokens, 0, latency, cost, event_name, span.trace_id, span.span_id)
                
                span.set_attributes({
                    # OTEL GenAI compliant attributes
                    "gen_ai.usage.input_tokens": input_tokens,
                    "gen_ai.usage.total_tokens": response.usage.total_tokens,
                    # Backward compatibility
                    "response.prompt_tokens": input_tokens,
                    "response.total_tokens": response.usage.total_tokens,
                    "latency_ms": latency
                })
                span.set_status(0)
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ✓ Embedding tracked: ${cost:.6f}")
                
                return response
            except Exception as e:
                span.set_status(1, str(e))
                raise
            finally:
                span.end()
        
        client.embeddings.create = tracked_embeddings_create
        
        # Track image generation
        original_images_generate = client.images.generate
        def tracked_images_generate(*args, **kwargs):
            model = kwargs.get("model", "dall-e-3")
            size = kwargs.get("size", "1024x1024")
            quality = kwargs.get("quality", "standard")
            n = kwargs.get("n", 1)  # Number of images
            
            # Extract event_name from agentbill_options if provided
            agentbill_options = kwargs.pop("agentbill_options", {})
            event_name = agentbill_options.get("event_name", "image_generation")
            
            if self.config.get("debug"):
                print(f"[AgentBill] Intercepting OpenAI images.generate: {model} {size} {quality}")
            
            start_time = time.time()
            span = self.tracer.start_span("openai.images.generate", {
                # OTEL GenAI compliant attributes
                "gen_ai.system": "openai",
                "gen_ai.request.model": model,
                "gen_ai.operation.name": "image_generation",
                # Backward compatibility
                "model": model,
                "provider": "openai",
                "size": size,
                "quality": quality
            })
            
            try:
                response = original_images_generate(*args, **kwargs)
                latency = (time.time() - start_time) * 1000
                
                # ✅ FIX: Track cost for images (was missing!)
                cost_per_image = calculate_image_cost(model, size, quality)
                total_cost = cost_per_image * n
                
                # Track as special case: input_tokens=n images, output_tokens=0
                self._track_usage(model, "openai", n, 0, latency, total_cost, event_name, span.trace_id, span.span_id)
                
                span.set_attributes({
                    "latency_ms": latency,
                    "num_images": n,
                    "cost": total_cost
                })
                span.set_status(0)
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ✓ Image generation tracked: ${total_cost:.4f} ({n} images)")
                
                return response
            except Exception as e:
                span.set_status(1, str(e))
                raise
            finally:
                span.end()
        
        client.images.generate = tracked_images_generate
        
        # Track audio transcription (Whisper)
        original_audio_transcriptions_create = client.audio.transcriptions.create
        def tracked_audio_transcriptions_create(*args, **kwargs):
            model = kwargs.get("model", "whisper-1")
            
            # Extract event_name from agentbill_options if provided
            agentbill_options = kwargs.pop("agentbill_options", {})
            event_name = agentbill_options.get("event_name", "audio_transcription")
            
            if self.config.get("debug"):
                print(f"[AgentBill] Intercepting OpenAI audio.transcriptions.create")
            
            start_time = time.time()
            span = self.tracer.start_span("openai.audio.transcriptions.create", {
                # OTEL GenAI compliant attributes
                "gen_ai.system": "openai",
                "gen_ai.request.model": model,
                "gen_ai.operation.name": "audio_transcription",
                # Backward compatibility
                "model": model,
                "provider": "openai"
            })
            
            try:
                response = original_audio_transcriptions_create(*args, **kwargs)
                latency = (time.time() - start_time) * 1000
                
                # ✅ FIX: Track cost for audio (was missing!)
                # Whisper pricing is per minute, estimate based on latency
                # Conservative estimate: latency * 2 (assumes real-time factor)
                duration_seconds = latency / 1000 * 2
                cost = calculate_audio_cost(model, duration_seconds=duration_seconds)
                
                # Track as special case
                self._track_usage(model, "openai", 1, 0, latency, cost, event_name, span.trace_id, span.span_id)
                
                span.set_attributes({
                    "latency_ms": latency,
                    "cost": cost
                })
                span.set_status(0)
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ✓ Audio transcription tracked: ${cost:.6f}")
                
                return response
            except Exception as e:
                span.set_status(1, str(e))
                raise
            finally:
                span.end()
        
        
        # Track text-to-speech
        original_audio_speech_create = client.audio.speech.create
        def tracked_audio_speech_create(*args, **kwargs):
            model = kwargs.get("model", "tts-1")
            input_text = kwargs.get("input", "")
            voice = kwargs.get("voice", "alloy")
            
            # Extract event_name from agentbill_options if provided
            agentbill_options = kwargs.pop("agentbill_options", {})
            event_name = agentbill_options.get("event_name", "text_to_speech")
            
            if self.config.get("debug"):
                print(f"[AgentBill] Intercepting OpenAI audio.speech.create: {model}")
            
            start_time = time.time()
            span = self.tracer.start_span("openai.audio.speech.create", {
                # OTEL GenAI compliant attributes
                "gen_ai.system": "openai",
                "gen_ai.request.model": model,
                "gen_ai.operation.name": "audio_speech",
                # Backward compatibility
                "model": model,
                "provider": "openai",
                "voice": voice
            })
            
            try:
                response = original_audio_speech_create(*args, **kwargs)
                latency = (time.time() - start_time) * 1000
                
                # ✅ FIX: Track cost for TTS (was missing!)
                chars = len(input_text)
                cost = calculate_audio_cost(model, chars=chars)
                
                # Track as special case: input_tokens=chars/4, output_tokens=0
                self._track_usage(model, "openai", chars // 4, 0, latency, cost, event_name, span.trace_id, span.span_id)
                
                span.set_attributes({
                    "latency_ms": latency,
                    "characters": chars,
                    "cost": cost
                })
                span.set_status(0)
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ✓ TTS tracked: ${cost:.6f} ({chars} chars)")
                
                return response
            except Exception as e:
                span.set_status(1, str(e))
                raise
            finally:
                span.end()
        
        client.audio.speech.create = tracked_audio_speech_create
        
        # Track moderations
        original_moderations_create = client.moderations.create
        def tracked_moderations_create(*args, **kwargs):
            if self.config.get("debug"):
                print(f"[AgentBill] Intercepting OpenAI moderations.create")
            
            start_time = time.time()
            span = self.tracer.start_span("openai.moderations.create", {
                # OTEL GenAI compliant attributes
                "gen_ai.system": "openai",
                "gen_ai.request.model": kwargs.get("model", "text-moderation-latest"),
                "gen_ai.operation.name": "moderation",
                # Backward compatibility
                "model": kwargs.get("model", "text-moderation-latest"),
                "provider": "openai"
            })
            
            try:
                response = original_moderations_create(*args, **kwargs)
                latency = (time.time() - start_time) * 1000
                span.set_attributes({"latency_ms": latency})
                span.set_status(0)
                return response
            except Exception as e:
                span.set_status(1, str(e))
                raise
            finally:
                span.end()
        
        client.moderations.create = tracked_moderations_create
        
        return client
    
    def wrap_anthropic(self, client: Any) -> Any:
        """Wrap Anthropic client with Cost Guard protection"""
        original_create = client.messages.create
        
        def tracked_create(*args, **kwargs):
            model = kwargs.get("model", "unknown")
            messages = kwargs.get("messages", [])
            max_tokens = kwargs.get("max_tokens", 1000)
            
            # Extract event_name from agentbill_options if provided (don't pass to Anthropic)
            agentbill_options = kwargs.pop("agentbill_options", {})
            event_name = agentbill_options.get("event_name", "ai_request")
            
            # Phase 1: Validate budget BEFORE API call
            validation = self._validate_request(model, messages, max_tokens)
            if not validation.get("allowed"):
                error_msg = validation.get("reason", "Budget limit reached")
                error = Exception(error_msg)
                error.code = "BUDGET_EXCEEDED"
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ❌ Request blocked: {error_msg}")
                raise error
            
            # Phase 2: Execute AI call
            start_time = time.time()
            span = self.tracer.start_span("anthropic.message", {
                "model": model,
                "provider": "anthropic"
            })
            
            try:
                response = original_create(*args, **kwargs)
                latency = (time.time() - start_time) * 1000
                
                # Phase 3: Track actual usage with trace correlation
                input_tokens = response.usage.input_tokens
                output_tokens = response.usage.output_tokens
                cost = self._estimate_cost(model, input_tokens, output_tokens, "anthropic")
                
                self._track_usage(model, "anthropic", input_tokens, output_tokens, latency, cost, event_name, span.trace_id, span.span_id)
                
                span.set_attributes({
                    # OTEL GenAI compliant attributes
                    "gen_ai.usage.input_tokens": input_tokens,
                    "gen_ai.usage.output_tokens": output_tokens,
                    "gen_ai.response.id": response.id if hasattr(response, 'id') else None,
                    "gen_ai.response.finish_reasons": [response.stop_reason] if hasattr(response, 'stop_reason') else [],
                    # Backward compatibility
                    "response.input_tokens": input_tokens,
                    "response.output_tokens": output_tokens,
                    "latency_ms": latency
                })
                span.set_status(0)
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ✓ Protected call completed: ${cost:.4f}")
                
                return response
            except Exception as e:
                span.set_status(1, str(e))
                raise
            finally:
                span.end()
        
        client.messages.create = tracked_create
        return client
    
    def wrap_bedrock(self, client: Any) -> Any:
        """Wrap AWS Bedrock client with Cost Guard protection"""
        original_invoke_model = client.invoke_model
        
        def tracked_invoke_model(*args, **kwargs):
            model = kwargs.get("modelId", "unknown")
            body_str = kwargs.get("body", "{}")
            
            # Parse body for messages
            import json
            try:
                body = json.loads(body_str)
                messages = body.get("messages", body.get("prompt", ""))
                max_tokens = body.get("max_tokens", body.get("max_tokens_to_sample", 1000))
            except:
                messages = ""
                max_tokens = 1000
            
            # Phase 1: Validate budget BEFORE API call
            validation = self._validate_request(model, messages, max_tokens)
            if not validation.get("allowed"):
                error_msg = validation.get("reason", "Budget limit reached")
                error = Exception(error_msg)
                error.code = "BUDGET_EXCEEDED"
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ❌ Request blocked: {error_msg}")
                raise error
            
            # Phase 2: Execute AI call
            if self.config.get("debug"):
                print(f"[AgentBill] Intercepting Bedrock invoke_model with model: {model}")
            
            start_time = time.time()
            span = self.tracer.start_span("bedrock.invoke_model", {
                # OTEL GenAI compliant attributes (using aws.bedrock as custom system)
                "gen_ai.system": "aws.bedrock",
                "gen_ai.request.model": model,
                "gen_ai.operation.name": "chat",
                # Backward compatibility
                "model": model,
                "provider": "bedrock"
            })
            
            try:
                response = original_invoke_model(*args, **kwargs)
                latency = (time.time() - start_time) * 1000
                
                # Parse response body for token usage
                body = json.loads(response['body'].read())
                
                # Handle different Bedrock model response formats
                input_tokens = 0
                output_tokens = 0
                if 'usage' in body:  # Claude models
                    input_tokens = body['usage'].get('input_tokens', 0)
                    output_tokens = body['usage'].get('output_tokens', 0)
                elif 'inputTextTokenCount' in body:  # Titan models
                    input_tokens = body.get('inputTextTokenCount', 0)
                    output_tokens = body['results'][0].get('tokenCount', 0) if 'results' in body else 0
                
                # Phase 3: Track actual usage with correlation
                cost = self._estimate_cost(model, input_tokens, output_tokens)
                self._track_usage(model, "bedrock", input_tokens, output_tokens, latency, cost, "ai_request", span.trace_id, span.span_id)
                
                span.set_attributes({
                    # OTEL GenAI compliant attributes
                    "gen_ai.usage.input_tokens": input_tokens,
                    "gen_ai.usage.output_tokens": output_tokens,
                    # Backward compatibility
                    "response.input_tokens": input_tokens,
                    "response.output_tokens": output_tokens,
                    "latency_ms": latency
                })
                span.set_status(0)
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ✓ Protected call completed: ${cost:.4f}")
                
                return response
            except Exception as e:
                span.set_status(1, str(e))
                raise
            finally:
                span.end()
        
        client.invoke_model = tracked_invoke_model
        return client
    
    def wrap_azure_openai(self, client: Any) -> Any:
        """Wrap Azure OpenAI client with Cost Guard protection"""
        original_create = client.chat.completions.create
        
        def tracked_create(*args, **kwargs):
            model = kwargs.get("model", "unknown")
            messages = kwargs.get("messages", [])
            max_tokens = kwargs.get("max_tokens", 1000)
            
            # Phase 1: Validate budget BEFORE API call
            validation = self._validate_request(model, messages, max_tokens)
            if not validation.get("allowed"):
                error_msg = validation.get("reason", "Budget limit reached")
                error = Exception(error_msg)
                error.code = "BUDGET_EXCEEDED"
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ❌ Request blocked: {error_msg}")
                raise error
            
            # Phase 2: Execute AI call
            if self.config.get("debug"):
                print(f"[AgentBill] Intercepting Azure OpenAI chat.completions.create")
            
            start_time = time.time()
            span = self.tracer.start_span("azure_openai.chat.completions", {
                # OTEL GenAI compliant attributes (using azure_openai as custom system)
                "gen_ai.system": "azure_openai",
                "gen_ai.request.model": model,
                "gen_ai.operation.name": "chat",
                # Backward compatibility
                "model": model,
                "provider": "azure_openai"
            })
            
            try:
                response = original_create(*args, **kwargs)
                latency = (time.time() - start_time) * 1000
                
                # Phase 3: Track actual usage
                input_tokens = response.usage.prompt_tokens
                output_tokens = response.usage.completion_tokens
                cost = self._estimate_cost(model, input_tokens, output_tokens)
                
                self._track_usage(model, "azure_openai", input_tokens, output_tokens, latency, cost, "ai_request", span.trace_id, span.span_id)
                
                span.set_attributes({
                    # OTEL GenAI compliant attributes
                    "gen_ai.usage.input_tokens": input_tokens,
                    "gen_ai.usage.output_tokens": output_tokens,
                    "gen_ai.usage.total_tokens": response.usage.total_tokens,
                    # Backward compatibility
                    "response.prompt_tokens": input_tokens,
                    "response.completion_tokens": output_tokens,
                    "response.total_tokens": response.usage.total_tokens,
                    "latency_ms": latency
                })
                span.set_status(0)
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ✓ Protected call completed: ${cost:.4f}")
                
                return response
            except Exception as e:
                span.set_status(1, str(e))
                raise
            finally:
                span.end()
        
        client.chat.completions.create = tracked_create
        return client
    
    def wrap_mistral(self, client: Any) -> Any:
        """Wrap Mistral AI client with Cost Guard protection"""
        original_create = client.chat.complete
        
        def tracked_create(*args, **kwargs):
            model = kwargs.get("model", "unknown")
            messages = kwargs.get("messages", [])
            max_tokens = kwargs.get("max_tokens", 1000)
            
            # Phase 1: Validate budget BEFORE API call
            validation = self._validate_request(model, messages, max_tokens)
            if not validation.get("allowed"):
                error_msg = validation.get("reason", "Budget limit reached")
                error = Exception(error_msg)
                error.code = "BUDGET_EXCEEDED"
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ❌ Request blocked: {error_msg}")
                raise error
            
            # Phase 2: Execute AI call
            if self.config.get("debug"):
                print(f"[AgentBill] Intercepting Mistral chat.complete with model: {model}")
            
            start_time = time.time()
            span = self.tracer.start_span("mistral.chat.complete", {
                # OTEL GenAI compliant attributes (using mistral as custom system)
                "gen_ai.system": "mistral",
                "gen_ai.request.model": model,
                "gen_ai.operation.name": "chat",
                # Backward compatibility
                "model": model,
                "provider": "mistral"
            })
            
            try:
                response = original_create(*args, **kwargs)
                latency = (time.time() - start_time) * 1000
                
                # Phase 3: Track actual usage
                input_tokens = response.usage.prompt_tokens
                output_tokens = response.usage.completion_tokens
                cost = self._estimate_cost(model, input_tokens, output_tokens)
                
                self._track_usage(model, "mistral", input_tokens, output_tokens, latency, cost, "ai_request", span.trace_id, span.span_id)
                
                span.set_attributes({
                    # OTEL GenAI compliant attributes
                    "gen_ai.usage.input_tokens": input_tokens,
                    "gen_ai.usage.output_tokens": output_tokens,
                    "gen_ai.usage.total_tokens": response.usage.total_tokens,
                    # Backward compatibility
                    "response.prompt_tokens": input_tokens,
                    "response.completion_tokens": output_tokens,
                    "response.total_tokens": response.usage.total_tokens,
                    "latency_ms": latency
                })
                span.set_status(0)
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ✓ Protected call completed: ${cost:.4f}")
                
                return response
            except Exception as e:
                span.set_status(1, str(e))
                raise
            finally:
                span.end()
        
        client.chat.complete = tracked_create
        return client
    
    def wrap_google_ai(self, client: Any) -> Any:
        """Wrap Google AI (Gemini) client with Cost Guard protection"""
        original_generate_content = client.generate_content
        
        def tracked_generate_content(*args, **kwargs):
            model = getattr(client, '_model_name', 'gemini-pro')
            content = args[0] if args else kwargs.get("contents", "")
            
            # Phase 1: Validate budget BEFORE API call
            validation = self._validate_request(model, content, 1000)
            if not validation.get("allowed"):
                error_msg = validation.get("reason", "Budget limit reached")
                error = Exception(error_msg)
                error.code = "BUDGET_EXCEEDED"
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ❌ Request blocked: {error_msg}")
                raise error
            
            # Phase 2: Execute AI call
            if self.config.get("debug"):
                print(f"[AgentBill] Intercepting Google AI generate_content")
            
            start_time = time.time()
            span = self.tracer.start_span("google_ai.generate_content", {
                # OTEL GenAI compliant attributes (using google_ai as custom system)
                "gen_ai.system": "google_ai",
                "gen_ai.request.model": model,
                "gen_ai.operation.name": "chat",
                # Backward compatibility
                "model": model,
                "provider": "google_ai"
            })
            
            try:
                response = original_generate_content(*args, **kwargs)
                latency = (time.time() - start_time) * 1000
                
                # Phase 3: Track actual usage
                input_tokens = 0
                output_tokens = 0
                if hasattr(response, 'usage_metadata'):
                    input_tokens = response.usage_metadata.prompt_token_count
                    output_tokens = response.usage_metadata.candidates_token_count
                
                cost = self._estimate_cost(model, input_tokens, output_tokens)
                self._track_usage(model, "google_ai", input_tokens, output_tokens, latency, cost, "ai_request", span.trace_id, span.span_id)
                
                if hasattr(response, 'usage_metadata'):
                    span.set_attributes({
                        # OTEL GenAI compliant attributes
                        "gen_ai.usage.input_tokens": input_tokens,
                        "gen_ai.usage.output_tokens": output_tokens,
                        "gen_ai.usage.total_tokens": response.usage_metadata.total_token_count,
                        # Backward compatibility
                        "response.prompt_tokens": input_tokens,
                        "response.completion_tokens": output_tokens,
                        "response.total_tokens": response.usage_metadata.total_token_count,
                        "latency_ms": latency
                    })
                
                span.set_status(0)
                
                if self.config.get("debug"):
                    print(f"[AgentBill Cost Guard] ✓ Protected call completed: ${cost:.4f}")
                
                return response
            except Exception as e:
                span.set_status(1, str(e))
                raise
            finally:
                span.end()
        
        client.generate_content = tracked_generate_content
        return client
    
    def wrap_perplexity(self, client: Any) -> Any:
        """Wrap Perplexity client with Cost Guard protection"""
        from .perplexity_wrapper import wrap_perplexity
        return wrap_perplexity(self, client)
    
    def wrap_ollama(self, client: Any) -> Any:
        """Wrap Ollama client with usage tracking (local, no API costs)"""
        from .ollama_wrapper import wrap_ollama
        return wrap_ollama(self, client)
    
    def track_signal(self, **kwargs):
        """
        Track a custom signal/event with comprehensive parameters
        
        Supports all 68 parameters including optional trace_id and span_id for OTEL correlation:
        - event_name (required)
        - trace_id (optional) - For correlating with OTEL spans for cost reconciliation
        - span_id (optional) - For correlating with OTEL spans for cost reconciliation
        - agent_external_id (auto-filled from config if not provided)
        - data_source, timestamp
        - customer_external_id, account_external_id, user_external_id, 
          order_external_id, session_id, conversation_id, thread_id
        - model, provider, prompt_hash, prompt_sample, response_sample, function_name, tool_name
        - prompt_tokens, completion_tokens, total_tokens, streaming_tokens, cached_tokens, reasoning_tokens
        - latency_ms, time_to_first_token, time_to_action_ms, queue_time_ms, processing_time_ms
        - revenue, cost, conversion_value, revenue_source
        - experiment_id, experiment_group, variant_id, ab_test_name
        - conversion_type, conversion_step, funnel_stage, goal_achieved
        - feedback_score, user_satisfaction, error_type, error_message, retry_count, success_rate
        - tags, category, priority, severity, compliance_flag, data_classification
        - product_id, feature_flag, environment, deployment_version, region, tenant_id
        - parent_span_id
        - custom_dimensions, metadata, data
        
        Example:
            # Basic tracking
            agentbill.track_signal(
                event_name="user_conversion",
                revenue=99.99,
                customer_external_id="cust_123"
            )
            
            # With OTEL correlation for cost reconciliation
            trace_context = agentbill.tracer.start_span("ai_completion")
            # ... make AI call ...
            agentbill.track_signal(
                event_name="ai_request",
                revenue=5.00,
                trace_id=trace_context.trace_id,  # Optional
                span_id=trace_context.span_id     # Optional
            )
        """
        import httpx
        import time
        
        if "event_name" not in kwargs:
            raise ValueError("event_name is required")
        
        url = f"{self.config.get('base_url', 'https://bgwyprqxtdreuutzpbgw.supabase.co')}/functions/v1/record-signals"
        
        # Build payload with all provided parameters
        payload = {k: v for k, v in kwargs.items() if v is not None}
        
        # Auto-fill agent_id or agent_external_id from config if not provided (REQUIRED by API)
        if "agent_external_id" not in payload and "agent_id" not in payload:
            if "agent_external_id" in self.config:
                payload["agent_external_id"] = self.config["agent_external_id"]
            elif "agent_id" in self.config:
                payload["agent_id"] = self.config["agent_id"]
            else:
                raise ValueError(
                    "agent_id or agent_external_id is required. Either pass it in track_signal() or set it in AgentBill.init() config. "
                    "Example: AgentBill.init({'api_key': '...', 'agent_id': 'uuid-here'}) or agent_external_id: 'my-agent-1'"
                )
        
        # Auto-fill customer_id or customer_external_id from config if not provided
        if "customer_id" not in payload and "customer_external_id" not in payload:
            if "customer_id" in self.config:
                # Check if it's a UUID format - send as customer_id, else customer_external_id
                import re
                customer_val = self.config["customer_id"]
                is_uuid = bool(re.match(r'^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$', customer_val, re.I))
                if is_uuid:
                    payload["customer_id"] = customer_val
                else:
                    payload["customer_external_id"] = customer_val
            elif "customer_external_id" in self.config:
                payload["customer_external_id"] = self.config["customer_external_id"]
        
        # Add timestamp if not provided
        if "timestamp" not in payload:
            payload["timestamp"] = time.time()
        
        try:
            with httpx.Client() as client:
                response = client.post(
                    url,
                    json=payload,
                    headers={
                        "X-API-Key": self.config['api_key'],
                        "Content-Type": "application/json"
                    },
                    timeout=10
                )
                if self.config.get("debug"):
                    trace_info = f" (trace: {kwargs.get('trace_id')})" if kwargs.get('trace_id') else ""
                    print(f"[AgentBill] Signal tracked: {kwargs.get('event_name')}{trace_info}")
                return response.status_code == 200
        except Exception as e:
            if self.config.get("debug"):
                print(f"[AgentBill] Failed to track signal: {e}")
            return False
    
    def track_conversion(self, event_type: str, event_value: float, signal_id: Optional[str] = None, 
                         session_id: Optional[str] = None, attribution_window_hours: int = 24,
                         currency: str = "USD", metadata: Optional[Dict] = None) -> Dict[str, Any]:
        """
        Track a conversion event for prompt profitability analysis
        Links conversions to AI prompts to calculate ROI per prompt
        
        Args:
            event_type: Type of conversion (e.g., 'purchase', 'signup', 'trial_start')
            event_value: Revenue amount from the conversion
            signal_id: Optional UUID linking to specific AI signal/prompt
            session_id: Optional session identifier
            attribution_window_hours: Time window for attribution (default: 24 hours)
            currency: Currency code (default: 'USD')
            metadata: Optional additional data
            
        Returns:
            Dict with success status and conversion_id
            
        Example:
            # Track a purchase conversion
            result = agentbill.track_conversion(
                event_type="purchase",
                event_value=99.99,
                currency="USD"
            )
            
            # Link conversion to specific AI prompt
            result = agentbill.track_conversion(
                event_type="trial_signup",
                event_value=29.99,
                signal_id="signal-uuid-from-prompt",
                session_id="session-123"
            )
        """
        import httpx
        
        url = f"{self.config.get('base_url', 'https://bgwyprqxtdreuutzpbgw.supabase.co')}/functions/v1/track-conversion"
        
        payload = {
            "api_key": self.config["api_key"],
            "customer_id": self.config.get("customer_id"),
            "event_type": event_type,
            "event_value": event_value,
            "signal_id": signal_id,
            "session_id": session_id,
            "attribution_window_hours": attribution_window_hours,
            "currency": currency,
            "metadata": metadata
        }
        
        try:
            with httpx.Client() as client:
                response = client.post(
                    url,
                    json=payload,
                    headers={"Content-Type": "application/json"},
                    timeout=10
                )
                data = response.json()
                
                if response.status_code != 200:
                    return {
                        "success": False,
                        "error": data.get("error", f"HTTP {response.status_code}")
                    }
                
                if self.config.get("debug"):
                    print(f"[AgentBill] Conversion tracked: {event_type} = ${event_value} (ID: {data.get('conversion_id')})")
                
                return {
                    "success": True,
                    "conversion_id": data.get("conversion_id")
                }
        except Exception as e:
            if self.config.get("debug"):
                print(f"[AgentBill] Failed to track conversion: {e}")
            return {
                "success": False,
                "error": str(e)
            }
    
    async def flush(self):
        """Flush pending telemetry data"""
        await self.tracer.flush()
