import requests
from functools import wraps
from typing import Any, Dict, List

class HanditTracker:
    def __init__(self):
        self.tracking_server_url = "https://handit-api-299768392189.us-central1.run.app/api/track"
        self.performance_server_url = "https://handit-api-299768392189.us-central1.run.app/api/performance"
        self.api_key = None
        self.urls_to_track = []

    def config(self, api_key, tracking_url=None):
        if not api_key:
            raise ValueError("API key is required for configuration.")
        self.api_key = api_key
        if tracking_url:
            self.tracking_server_url = tracking_url
        print("Library configured successfully with API key.")

    def update_tracked_urls(self):
        if not self.api_key:
            raise ValueError("API key not set. Call the config method with your API key.")
        
        headers = {"Authorization": f"Bearer {self.api_key}"}
        try:
            response = requests.get(f"{self.tracking_server_url}/urls-to-track", headers=headers)
            response.raise_for_status()
            self.urls_to_track = response.json()
        except requests.RequestException as e:
            print(f"Error fetching URLs to track: {e}")

    def intercept_requests(self, func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            self.update_tracked_urls()
            url = args[0]
            matching_url = next((u for u in self.urls_to_track if u["url"] in url), None)
            if matching_url:
                model_id = matching_url["id"]
                request_body = kwargs.get("json", kwargs.get("data"))
                
                try:
                    response = func(*args, **kwargs)
                    response_body = response.json()
                    self._send_tracked_data(model_id, request_body, response_body)
                    return response
                except Exception as e:
                    print(f"Error tracking request: {e}")
                    raise
            else:
                return func(*args, **kwargs)
        return wrapper

    def capture_model(self, model_id, request_body, response_body):
        try:
            self._send_tracked_data(model_id, request_body, response_body)
        except Exception as e:
            print(f"Error in manual tracking: {e}")

    def _send_tracked_data(self, model_id, request_body, response_body):
        headers = {"Authorization": f"Bearer {self.api_key}"}
        payload = {
            "input": request_body,
            "output": response_body,
            "modelId": model_id,
            "parameters": {}
        }
        try:
            response = requests.post(self.tracking_server_url, json=payload, headers=headers)
            response.raise_for_status()
        except requests.RequestException as e:
            print(f"Error sending tracked data: {e}")

    def fetch_optimized_prompt(self, model_id):
        """
        Fetches the most optimized prompt for a given model ID.
        
        Args:
            model_id (str): The ID of the model to fetch the optimized prompt for.
            
        Returns:
            dict: The optimized prompt data
            
        Raises:
            ValueError: If API key is not configured
            requests.RequestException: If the API request fails
        """
        if not self.api_key:
            raise ValueError("API key not set. Call the config method with your API key.")
        
        headers = {"Authorization": f"Bearer {self.api_key}"}
        url = f"{self.performance_server_url}/model/{model_id}/optimized-prompt"
        
        try:
            response = requests.get(url, headers=headers)
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            print(f"Error fetching optimized prompt: {e}")
            raise

    def track_model(self, model, model_id):
        """
        Creates a tracked version of an AI model (LLM, embedding model, etc.)
        
        Args:
            model: The model to track (ChatOpenAI, OpenAIEmbeddings, etc.)
            model_id (str): The ID of the model to track
            
        Returns:
            Modified model with tracking capabilities
        """
        original_call = model.__call__
        
        async def tracked_call(*args, **kwargs):
            try:
                # Track model input
                self._send_tracked_data(
                    model_id=model_id,
                    request_body={
                        "type": "model_call",
                        "model_type": model.__class__.__name__,
                        "messages": kwargs.get("messages", []),
                        "prompt": kwargs.get("prompt", ""),
                        "parameters": {
                            k: v for k, v in kwargs.items()
                            if k not in ["messages", "prompt"]
                        }
                    },
                    response_body=None
                )
                
                # Execute model
                result = await original_call(*args, **kwargs)
                
                # Track model output
                self._send_tracked_data(
                    model_id=model_id,
                    request_body={
                        "type": "model_response",
                        "model_type": model.__class__.__name__
                    },
                    response_body={
                        "response": str(result)
                    }
                )
                
                return result
            except Exception as e:
                print(f"Error tracking model: {e}")
                raise
        
        model.__call__ = tracked_call
        return model

    def track_tool(self, tool, tool_id):
        """
        Creates a tracked version of a tool (RAG, API call, function, etc.)
        
        Args:
            tool: The tool to track
            tool_id (str): The ID of the tool to track
            
        Returns:
            Modified tool with tracking capabilities
        """
        original_call = tool.__call__ if hasattr(tool, '__call__') else tool
        
        async def tracked_call(*args, **kwargs):
            try:
                # Track tool input
                self._send_tracked_data(
                    model_id=tool_id,  # using model_id field for tool_id
                    request_body={
                        "type": "tool_call",
                        "tool_type": tool.__class__.__name__,
                        "inputs": {
                            "args": args,
                            "kwargs": kwargs
                        }
                    },
                    response_body=None
                )
                
                # Execute tool
                result = await original_call(*args, **kwargs)
                
                # Track tool output
                self._send_tracked_data(
                    model_id=tool_id,
                    request_body={
                        "type": "tool_response",
                        "tool_type": tool.__class__.__name__
                    },
                    response_body={
                        "result": result
                    }
                )
                
                return result
            except Exception as e:
                print(f"Error tracking tool: {e}")
                raise
        
        if hasattr(tool, '__call__'):
            tool.__call__ = tracked_call
            return tool
        else:
            return tracked_call
