import os
import time
import uuid
import json
import requests
from functools import wraps
import threading

class Leap360Tracer:
    def __init__(self, api_key=None, service_name="python-sdk"):
        self.api_key = api_key or self._load_api_key()
        self.service_name = service_name
        self.api_url = os.environ.get("LEAP360_API_URL", "https://b0vyk178b7.execute-api.us-east-1.amazonaws.com")
        
    def _load_api_key(self):
        """Loads the Leap360 API key from environment variables or local CLI config."""
        key = os.environ.get("LEAP360_API_KEY")
        if key:
            return key
        config_path = os.path.expanduser("~/.leap360/config.json")
        if os.path.exists(config_path):
            try:
                with open(config_path, "r", encoding="utf-8") as f:
                    return json.load(f).get("apiKey")
            except Exception:
                pass
        return None

    def _send_trace_async(self, payload):
        if not self.api_key:
            return  # Silent drop if not authenticated

        def _post():
            try:
                requests.post(
                    f"{self.api_url}/v1/traces",
                    headers={"Content-Type": "application/json", "x-leap360-key": self.api_key},
                    json=payload,
                    timeout=5
                )
            except Exception:
                pass # Fail silently for telemetry

        # Send telemetry in background thread so we don't block the AI generation loop
        threading.Thread(target=_post, daemon=True).start()

    def wrap(self, client):
        """
        Wraps an LLM client (e.g., openai.OpenAI) natively catching all API interactions.
        """
        client_type_str = str(type(client)).lower()
        if "openai" in client_type_str and hasattr(client, "chat"):
            return self._wrap_openai(client)
        # Fallback/passthrough if the client is not natively hooked yet
        return client
        
    def _wrap_openai(self, client):
        original_create = client.chat.completions.create
        
        @wraps(original_create)
        def patched_create(*args, **kwargs):
            start_time = time.time()
            error_str = None
            response = None
            try:
                response = original_create(*args, **kwargs)
                return response
            except Exception as e:
                error_str = str(e)
                raise
            finally:
                latency_ms = int((time.time() - start_time) * 1000)
                
                # Extract usage metadata
                usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
                if response and hasattr(response, "usage") and response.usage:
                    usage = {
                        "prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
                        "completion_tokens": getattr(response.usage, "completion_tokens", 0),
                        "total_tokens": getattr(response.usage, "total_tokens", 0),
                    }
                    
                model = kwargs.get("model", "unknown-openai-model")
                req_data = {"messages": kwargs.get("messages", [])}
                
                # Format payload compliant with the Leap360 Cloud ingestion endpoint
                payload = {
                    "service_name": self.service_name,
                    "run_id": str(uuid.uuid4()),
                    "provider": "openai",
                    "model": model,
                    "request": req_data,
                    "response": response.model_dump() if hasattr(response, "model_dump") else str(response),
                    "latency_ms": latency_ms,
                    "usage": usage,
                    "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                }
                if error_str:
                    payload["error"] = error_str
                    
                self._send_trace_async(payload)
                
        client.chat.completions.create = patched_create
        return client
