import requests
import time
import os
from typing import Optional, Dict, Any
from .core import Netra, HealthPulse

# --- SHARED HELPER FUNCTION ---
def _perform_health_check(agent_id: str, url: str, headers: Dict[str, str], timeout: int = 5) -> HealthPulse:
    """
    Internal helper to ping a URL and return a standardized HealthPulse.
    Used by both CrewAI and LangGraph adapters to avoid code duplication.
    """
    start = time.time()
    status = "down"
    score = 0
    
    try:
        # Perform the actual network request
        response = requests.get(url, headers=headers, timeout=timeout)
        
        # --- STRICT STATUS CHECK ---
        # 200 = Healthy (Key worked, Server up)
        # 401 = DOWN (Key Failed - Critical Logic)
        # 404 = Degraded (Server up, endpoint might be slightly wrong but alive)
        
        if response.status_code == 200:
            status = "healthy"
            score = 100
        elif response.status_code == 401:
            status = "down" 
            score = 0
            print(f"⚠️ SAKSH WARNING: Agent {agent_id} Failed Auth (401)")
        else:
            status = "degraded"
            score = 50
            
    except requests.exceptions.Timeout:
        status = "degraded" # It's alive but slow
        score = 40
        print(f"⚠️ SAKSH WARNING: Agent {agent_id} Timed Out")
    except Exception:
        status = "down"
        score = 0
        
    # Calculate Latency
    latency = (time.time() - start) * 1000
    
    # Adjust Score based on Latency (Penalize Slowness)
    if status == "healthy":
        if latency > 1500: score -= 20
        if latency > 3000: score -= 40
        
    return HealthPulse(
        agent_id=agent_id,
        status=status,
        latency_ms=round(latency, 2),
        vitality_score=max(0, score) # Ensure score never goes below 0
    )


# --- 1. CREWAI ADAPTER ---
class CrewAINetra(Netra):
    """
    Specialized for CrewAI Agents.
    Automatically extracts: agent.llm.base_url and agent.llm.api_key
    """
    def __init__(self, crewai_agent, agent_id: str):
        super().__init__(agent_id)
        self.agent = crewai_agent
        
        # 1. Logic to extract URL
        # CrewAI often hides the base_url inside the llm object
        if hasattr(crewai_agent.llm, 'base_url') and crewai_agent.llm.base_url:
            self.target_url = crewai_agent.llm.base_url
            self.is_custom = True
        else:
            self.target_url = "https://api.openai.com/v1"
            self.is_custom = False
            
        # 2. Logic to extract API Key
        self.api_key = getattr(crewai_agent.llm, 'api_key', None)
        
        # Fallback: Check Environment Variable if key is missing in object
        if not self.api_key and not self.is_custom:
             self.api_key = os.getenv("OPENAI_API_KEY")

    def gaze(self) -> HealthPulse:
        # Construct the Check URL
        if self.is_custom:
            base = self.target_url.rstrip("/")
            if base.endswith("/v1"):
                check_url = f"{base}/models"
            else:
                check_url = f"{base}/v1/models"
        else:
            check_url = "https://api.openai.com/v1/models"

        # Construct Headers
        headers = {}
        if self.api_key and self.api_key != "na":
            headers["Authorization"] = f"Bearer {self.api_key}"

        # Delegate to shared helper
        return _perform_health_check(self.agent_id, check_url, headers)


# --- 2. LANGGRAPH ADAPTER ---
class LangGraphNetra(Netra):
    """
    Specialized for LangGraph Compiled Graphs.
    Since Graphs are abstract, we monitor the underlying LLM provider they rely on.
    """
    def __init__(self, graph_app, agent_id: str, target_url="https://api.openai.com/v1", api_key=None):
        super().__init__(agent_id)
        self.graph = graph_app
        
        # LangGraph allows flexible backends, so we accept overrides.
        # Defaults to OpenAI if not specified.
        self.target_url = target_url
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")

    def gaze(self) -> HealthPulse:
        # Construct URL
        base = self.target_url.rstrip("/")
        if base.endswith("/v1"):
            check_url = f"{base}/models"
        else:
            check_url = f"{base}/v1/models"
            
        # Construct Headers
        headers = {}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        # Delegate to shared helper
        return _perform_health_check(self.agent_id, check_url, headers)