"""
LLM Evaluator – analyzes BESS-RL evaluation results using the Hugging Face Router.
Includes a heuristic fallback for when API quotas are exceeded.
All models are routed via the Hugging Face OpenAI-compatible API.
"""
import os
import json
import re
import requests
import openai
from dotenv import load_dotenv

# Load environment variables from .env
load_dotenv()

# Default Model IDs for the HF Router
DEFAULT_QWEN_MODEL = os.getenv("HF_QWEN_MODEL", "Qwen/Qwen2.5-72B-Instruct")
DEFAULT_GPT_MODEL = os.getenv("HF_GPT_MODEL", "openai/gpt-oss-20b")
DEFAULT_FALLBACK_MODEL = "meta-llama/Llama-3.3-70B-Instruct"

TASK_DESCRIPTIONS = {
    "easy":   "Energy Arbitrage only (buy cheap, sell expensive)",
    "medium": "Energy Arbitrage + Frequency Regulation",
    "hard":   "Energy Arbitrage + Frequency Regulation + Peak Shaving",
}

_FALLBACK_BASE = {
    "available": False,
    "verdict": "N/A",
    "score": 0.0,
    "reward_score": 0.0,
    "summary": "",
    "strengths": [],
    "weaknesses": [],
    "recommendations": [],
    "confidence": "N/A",
    "detailed_analysis": "",
    "error": None,
    "is_heuristic": False,
    "provider": "UNKNOWN"
}

def _extract_json(text: str) -> dict:
    """Robustly extracts JSON from a string, handling markdown or surrounding text."""
    text = text.strip()
    parsed = {}
    
    # Try direct parse first
    try:
        parsed = json.loads(text)
    except:
        # Try finding JSON block with regex
        json_pattern = r'\{.*\}'
        matches = re.findall(json_pattern, text, re.DOTALL)
        for m in reversed(matches):
            try:
                parsed = json.loads(m)
                break
            except:
                continue
                
    # If no valid JSON found, use a fallback structure
    if not isinstance(parsed, dict) or not parsed:
        parsed = {
            "verdict": "Unknown",
            "score": 0.0,
            "reward_score": 0.0,
            "summary": "Parsing failed.",
            "detailed_analysis": text[:500] if text else "No content"
        }
            
    # Manual regex for key floating point fields to be safe
    for field in ["score", "reward_score"]:
        if field not in parsed or parsed[field] == 0.0:
            match = re.search(fr'"{field}"\s*:\s*([\d\.]+)', text)
            if match:
                try: parsed[field] = float(match.group(1))
                except: pass
        if field in parsed and isinstance(parsed[field], (int, float)):
            parsed[field] = max(0.001, min(0.999, float(parsed[field])))

    return parsed

def _get_heuristic_analysis(data: dict) -> dict:
    """Provides a rule-based assessment when LLMs are unavailable."""
    scores = data.get("scores", {})
    overall = scores.get("overall", 0)
    task = data.get("task", "hard")
    
    if overall >= 0.85: verdict = "Excellent"
    elif overall >= 0.70: verdict = "Good"
    elif overall >= 0.50: verdict = "Needs Improvement"
    else: verdict = "Poor"
    
    strengths, weaknesses, recommendations = [], [], []
    if scores.get("reward", 0) > 0.8: strengths.append("Strong reward optimization")
    else: weaknesses.append("Sub-optimal profit generation")
    
    if scores.get("soc_readiness", 0) > 0.8: strengths.append("Excellent peak-hour SOC management")
    else: recommendations.append("Improve SOC buffer management before evening peaks")
    
    if scores.get("ps_adherence", 0) < 0.9 and task == "hard":
        weaknesses.append("Frequent peak shaving violations")
        recommendations.append("Increase penalty factor for grid load violations in reward function")
        
    if scores.get("cycle_discipline", 0) < 0.6:
        weaknesses.append("Aggressive battery cycling")
        recommendations.append("Increase degradation cost coefficient to preserve battery health")

    h_score = (
        scores.get("reward", 0) * 0.4 +
        scores.get("soc_readiness", 0) * 0.2 +
        scores.get("ps_adherence", 0) * 0.3 +
        scores.get("cycle_discipline", 0) * 0.1
    )

    return {
        **_FALLBACK_BASE,
        "available": True,
        "is_heuristic": True,
        "provider": "LOCAL",
        "verdict": verdict,
        "score": max(0.001, min(0.999, round(h_score, 3))),
        "reward_score": max(0.001, min(0.999, round(scores.get("reward", 0), 3))),
        "summary": f"Heuristic assessment: The agent shows a {verdict.lower()} grasp of the {task} task.",
        "strengths": strengths or ["Stable baseline performance"],
        "weaknesses": weaknesses or ["Minor inefficiencies"],
        "recommendations": recommendations or ["Continue training"],
        "confidence": "Medium (Rule-based)",
        "detailed_analysis": "Local rule-based assessment."
    }

def _build_prompt(data: dict) -> str:
    task = data.get("task", "unknown")
    scores = data.get("scores", {})
    return f"""Analyze BESS RL Agent Performance:
Task: {task}
Mean Reward: {data.get('reward_mean', 0):.1f}
Simulated Reward Score: {scores.get('reward', 0):.3f}
Simulated Overall Score: {scores.get('overall', 0):.3f}

Return ONLY this JSON format:
{{
  "verdict": "Excellent|Good|Needs Improvement|Poor",
  "score": <float 0-1 overall assessment>,
  "reward_score": <float 0-1 specifically for profit performance>,
  "summary": "...",
  "strengths": ["..."],
  "weaknesses": ["..."],
  "recommendations": ["..."],
  "confidence": "High|Medium|Low",
  "detailed_analysis": "..."
}}
"""

def _call_router(client: openai.OpenAI, model_id: str, prompt: str) -> dict:
    """Internal helper to make the chat completion call."""
    response = client.chat.completions.create(
        model=model_id,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.3
    )
    content = response.choices[0].message.content
    parsed = _extract_json(content)
    # Ensure all required fields are present to prevent frontend crashes
    for key, val in _FALLBACK_BASE.items():
        if key not in parsed and key != "available":
            parsed[key] = val
    return {**parsed, "available": True, "provider": f"HF:{model_id}"}


def _get_router_analysis(data: dict, model_id: str) -> dict:
    """Calls the Hugging Face Router with automatic fallback support."""
    api_token = os.getenv("HF_TOKEN") or os.getenv("PowerGrid") or os.getenv("HUGGING_FACE_HUB_TOKEN")
    if not api_token:
        raise ValueError("Hugging Face token not found (expected HF_TOKEN).")
    
    client = openai.OpenAI(
        base_url="https://router.huggingface.co/v1",
        api_key=api_token
    )
    
    prompt = _build_prompt(data)
    
    try:
        return _call_router(client, model_id, prompt)
    except Exception as e:
        error_msg = str(e)
        if "model_not_found" in error_msg or "does not exist" in error_msg or "404" in error_msg:
            try:
                result = _call_router(client, DEFAULT_FALLBACK_MODEL, prompt)
                result["summary"] = f"NOTE: Primary model {model_id} unavailable. Falling back. " + result.get("summary", "")
                return result
            except:
                pass
        raise ValueError(f"Failed to analyze: {error_msg}")

def get_llm_analysis(data: dict, provider: str = "QWEN", model_name: str = None) -> dict:
    """Main entry point for technical LLM analysis via Hugging Face Router."""
    try:
        if provider == "QWEN":
            model_id = model_name or DEFAULT_QWEN_MODEL
        else:
            model_id = model_name or DEFAULT_GPT_MODEL
            
        return _get_router_analysis(data, model_id)
        
    except Exception as e:
        error_msg = str(e)
        return {
            **_FALLBACK_BASE,
            "available": False,
            "summary": f"Evaluation Failed: {error_msg}",
            "error": error_msg,
            "detailed_analysis": f"The requested analysis could not be completed.\n\nReason: {error_msg}"
        }