"""
Automated Model Benchmark: Reasoning, Tool Use & L3 Batching
------------------------------------------------------------
1. Loops through specified models.
2. Updates a persistent Markdown report.
3. BEHAVIOR: One table per Provider.
4. GUARANTEES ID UNIQUENESS: Merges new results based on Endpoint ID.
5. TESTS: Inference -> Single Tool -> Parallel Batch Tools.
"""

import json
import os
import re
import time
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from config_benchmark import config
from dotenv import load_dotenv

# --- SDK Event Imports ---
from projectdavid import PlanEvent  # [NEW] For L3 Strategy Visibility
from projectdavid import (
    ContentEvent,
    DecisionEvent,
    Entity,
    ReasoningEvent,
    StatusEvent,
    ToolCallRequestEvent,
)

# ------------------------------------------------------------------
# 0. Setup & Config
# ------------------------------------------------------------------
root_dir = Path(__file__).resolve().parents[2]
load_dotenv()

# ANSI Colors
CYAN = "\033[96m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
RED = "\033[91m"
MAGENTA = "\033[95m"
BLUE = "\033[94m"  # Planning Color
RESET = "\033[0m"

# --- Load External Config ---
CONFIG_FILE_NAME = "benchmark_config.json"
CONFIG_PATH = root_dir / CONFIG_FILE_NAME

# --- Global Constants ---
BASE_URL = config.get("base_url") or os.getenv("BASE_URL", "http://localhost:9000")
ENTITIES_API_KEY = config.get("entities_api_key") or os.getenv("ENTITIES_API_KEY")
USER_ID = config.get("entities_user_id") or os.getenv("ENTITIES_USER_ID")
report_name = config.get("report_file_name", "model_compatibility_report.md")
REPORT_FILE = root_dir / report_name


def get_api_key_for_provider(provider: str) -> str:
    """Dynamic key resolver."""
    p = provider.lower()
    if "together" in p:
        return config.get("together_api_key") or os.getenv("TOGETHER_API_KEY", "")
    elif "hyperbolic" in p:
        return config.get("hyperbolic_api_key") or os.getenv("HYPERBOLIC_API_KEY", "")
    elif "openai" in p:
        return config.get("openai_api_key") or os.getenv("OPENAI_API_KEY", "")
    return ""


# ------------------------------------------------------------------
# 1. Tool Logic (Registry)
# ------------------------------------------------------------------
def get_flight_times(tool_name: str, arguments: dict) -> str:
    """Mock flight-time lookup tool."""
    return json.dumps(
        {
            "status": "success",
            "duration": "4h 30m",
            "departure_time": "10:00 AM PST",
            "arrival_time": "06:30 PM EST",
            "route": f"{arguments.get('departure', 'UNK')} -> {arguments.get('arrival', 'UNK')}",
        }
    )


def get_weather(tool_name: str, arguments: dict) -> str:
    """Mock weather tool for Batch testing."""
    loc = arguments.get("location", "Unknown")
    return json.dumps(
        {"status": "success", "location": loc, "temp": "15C", "condition": "Cloudy"}
    )


# Registry for Dynamic Dispatch
TOOL_REGISTRY = {
    "get_flight_times": get_flight_times,
    "get_weather": get_weather,
}


# ------------------------------------------------------------------
# 2. Report Manager
# ------------------------------------------------------------------
class ReportManager:
    @staticmethod
    def update_report(new_results: List[Dict]):
        """
        Updates the MD file with Parallel Call column.
        """
        file_path = Path(REPORT_FILE)
        rows_by_provider = defaultdict(dict)

        # 1. Read Existing File
        if file_path.exists():
            try:
                with open(file_path, "r", encoding="utf-8") as f:
                    for line in f:
                        clean_line = line.strip()
                        if clean_line.startswith("| **"):
                            parts = [p.strip() for p in clean_line.split("|")]
                            # Ensure we capture enough columns for the new format
                            if len(parts) >= 4:
                                provider = parts[2].strip()
                                match = re.search(r"`(.*?)`", parts[3])
                                if match:
                                    eid = match.group(1).strip()
                                    rows_by_provider[provider][eid] = clean_line
            except Exception as e:
                print(f"{RED}[!] Error parsing report: {e}{RESET}")

        # 2. Merge New Results
        for res in new_results:
            inf_icon = "✅" if res["inference_ok"] else "❌"
            reas_icon = "🧠" if res["reasoning_detected"] else "—"
            tool_icon = "✅" if res["tool_call_ok"] else "❌"

            # 🚀 = Batch Success (2+ tools), — = Not tested/Failed
            batch_icon = "🚀" if res.get("batch_ok") else "—"
            if not res["tool_call_ok"]:
                batch_icon = "—"

            telem_icon = "📡" if res["call_telemetry"] else "—"

            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
            note = res["error_msg"] if res["error_msg"] else "OK"

            new_row = (
                f"| **{res['name']}** | {res['provider']} | `{res['id']}` | "
                f"{inf_icon} | {reas_icon} | {tool_icon} | {batch_icon} | {telem_icon} | {timestamp} | {note} |"
            )
            rows_by_provider[res["provider"]][res["id"]] = new_row

        # 3. Write Report
        table_header_row = "| Model Name | Provider | Endpoint ID | Inference | Reasoning | Tools | Parallel | Telemetry | Last Run | Notes |"
        table_align_row = "| :--- | :--- | :--- | :---: | :---: | :---: | :---: | :---: | :--- | :--- |"

        main_title = [
            "# 🧪 Model Compatibility Report",
            f"**Last Update:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
            "",
        ]

        try:
            with open(file_path, "w", encoding="utf-8") as f:
                f.write("\n".join(main_title))
                for provider in sorted(rows_by_provider.keys()):
                    f.write(f"\n## 🟢 Provider: {provider}\n")
                    f.write(f"{table_header_row}\n{table_align_row}\n")
                    rows = sorted(
                        rows_by_provider[provider].values(),
                        key=lambda x: x.split("|")[1].strip().lower(),
                    )
                    f.write("\n".join(rows) + "\n")
            print(f"{GREEN}[✓] Report updated: {file_path.name}{RESET}")
        except Exception as e:
            print(f"{RED}[!] Write failed: {e}{RESET}")


# ------------------------------------------------------------------
# 3. Test Logic
# ------------------------------------------------------------------
class ModelTester:
    def __init__(self, model_label: str, config: Dict):
        self.model_label = model_label
        self.model_id = config["id"]
        self.provider = config["provider"]
        self.client = Entity(base_url=BASE_URL, api_key=ENTITIES_API_KEY)
        self.assistant_id: Optional[str] = None
        self.user_id = USER_ID
        self.provider_api_key = get_api_key_for_provider(self.provider)

        # Bind clients for recursive SDK handling
        if hasattr(self.client, "synchronous_inference_stream"):
            self.client.synchronous_inference_stream.bind_clients(
                self.client.runs,
                self.client.actions,
                self.client.messages,
                self.client.assistants,
            )

    def _classify_error(self, error_text: str) -> str:
        text = str(error_text).lower()
        if "404" in text or "not found" in text:
            return "💀 DEAD / UNAVAILABLE"
        if "401" in text or "unauthorized" in text:
            return "⛔ AUTH ERROR"
        if "429" in text:
            return "⏳ RATE LIMIT"
        if "500" in text:
            return "🔥 SERVER ERROR"
        clean_err = str(error_text).replace("\n", " ")
        return f"Error: {clean_err[:30]}..."

    def setup(self) -> Tuple[bool, str]:
        if not self.provider_api_key:
            return False, f"Missing API Key for {self.provider}"

        try:
            print(f"{YELLOW}[*] [{self.model_label}] Creating Assistant...{RESET}")
            assistant = self.client.assistants.create_assistant(
                name=f"Bench_{self.model_label}",
                instructions="You are a helpful AI assistant. Use tools when needed.",
                model=self.model_id,
                tools=[
                    {
                        "type": "function",
                        "function": {
                            "name": "get_flight_times",
                            "description": "Get flight times",
                            "parameters": {
                                "type": "object",
                                "properties": {
                                    "departure": {"type": "string"},
                                    "arrival": {"type": "string"},
                                },
                                "required": ["departure", "arrival"],
                            },
                        },
                    },
                    {
                        "type": "function",
                        "function": {
                            "name": "get_weather",
                            "description": "Get current weather for a city",
                            "parameters": {
                                "type": "object",
                                "properties": {"location": {"type": "string"}},
                                "required": ["location"],
                            },
                        },
                    },
                ],
            )
            self.assistant_id = assistant.id
            return True, None
        except Exception as e:
            return False, self._classify_error(str(e))

    def run_benchmark(self) -> Dict:
        result = {
            "name": self.model_label,
            "id": self.model_id,
            "provider": self.provider,
            "inference_ok": False,
            "reasoning_detected": False,
            "tool_call_ok": False,
            "batch_ok": False,
            "call_telemetry": False,
            "error_msg": "",
        }

        # --- Stage 1: Inference ---
        print(f"\n{YELLOW}--- Stage 1: Inference ({self.model_label}) ---{RESET}")
        try:
            thread = self.client.threads.create_thread(participant_ids=[self.user_id])
            message = self.client.messages.create_message(
                thread_id=thread.id,
                role="user",
                content="Calculate Fibonacci recursively. Explain why iterative is better.",
                assistant_id=self.assistant_id,
            )
            run = self.client.runs.create_run(
                assistant_id=self.assistant_id, thread_id=thread.id
            )
            stats = self._stream_and_analyze(
                thread.id, message.id, run.id, timeout=180.0
            )

            result["inference_ok"] = stats["inference_ok"]
            result["reasoning_detected"] = stats["reasoning_detected"]
            if stats["error"]:
                result["error_msg"] = self._classify_error(stats["detailed_error"])
                return result

        except Exception as e:
            result["error_msg"] = self._classify_error(str(e))
            return result

        # --- Stage 2: Single Tool ---
        if result["inference_ok"]:
            print(f"\n{YELLOW}--- Stage 2: Single Tool ({self.model_label}) ---{RESET}")
            try:
                thread = self.client.threads.create_thread(
                    participant_ids=[self.user_id]
                )
                message = self.client.messages.create_message(
                    thread_id=thread.id,
                    role="user",
                    content="Fetch flight times between LAX and JFK.",
                    assistant_id=self.assistant_id,
                )
                run = self.client.runs.create_run(
                    assistant_id=self.assistant_id, thread_id=thread.id
                )
                stats = self._stream_and_analyze(
                    thread.id, message.id, run.id, timeout=120.0
                )

                if stats["decision_detected"]:
                    result["call_telemetry"] = True

                if stats["tool_executed_count"] > 0 and stats["inference_ok"]:
                    result["tool_call_ok"] = True
                    print(f"{GREEN}[✓] Single Tool Verified.{RESET}")
                else:
                    if stats["error"]:
                        result["error_msg"] = self._classify_error(
                            stats["detailed_error"]
                        )
                    elif stats["tool_executed_count"] == 0:
                        print(f"{RED}[!] Tool call not triggered.{RESET}")

            except Exception as e:
                result["error_msg"] = self._classify_error(str(e))

        # --- Stage 3: Batch/Parallel Tools (Level 3) ---
        if result["tool_call_ok"]:
            print(
                f"\n{MAGENTA}--- Stage 3: Parallel Batch ({self.model_label}) ---{RESET}"
            )
            try:
                thread = self.client.threads.create_thread(
                    participant_ids=[self.user_id]
                )
                # Prompt requires TWO tools
                message = self.client.messages.create_message(
                    thread_id=thread.id,
                    role="user",
                    content="I need two things: check the weather in London, AND find flight times from NYC to Paris.",
                    assistant_id=self.assistant_id,
                )
                run = self.client.runs.create_run(
                    assistant_id=self.assistant_id, thread_id=thread.id
                )
                stats = self._stream_and_analyze(
                    thread.id, message.id, run.id, timeout=120.0
                )

                # Check if we executed at least 2 distinct tools or 2 calls total
                if stats["tool_executed_count"] >= 2 and stats["inference_ok"]:
                    result["batch_ok"] = True
                    print(
                        f"{GREEN}[✓] Batch Verified: {stats['tool_executed_count']} tools executed.{RESET}"
                    )
                else:
                    print(
                        f"{RED}[!] Batch Failed. Tools executed: {stats['tool_executed_count']}{RESET}"
                    )

            except Exception as e:
                result["error_msg"] = f"Batch Error: {str(e)[:20]}"

        return result

    def _stream_and_analyze(self, thread_id, message_id, run_id, timeout) -> Dict:
        """Unified Event Loop for all stages."""
        stats = {
            "inference_ok": False,
            "reasoning_detected": False,
            "tool_executed_count": 0,
            "decision_detected": False,
            "error": False,
            "detailed_error": None,
        }

        stream = self.client.synchronous_inference_stream
        stream.setup(
            thread_id=thread_id,
            assistant_id=self.assistant_id,
            message_id=message_id,
            run_id=run_id,
            api_key=self.provider_api_key,
        )

        current_mode = None
        seen_tools = set()

        try:
            for event in stream.stream_events(
                provider=self.provider, model=self.model_id, timeout_per_chunk=timeout
            ):

                if isinstance(event, ReasoningEvent):
                    if current_mode != "reasoning":
                        print(f"\n{CYAN}🤔 [THOUGHT PROCESS]{RESET}\n", end="")
                        current_mode = "reasoning"
                        stats["reasoning_detected"] = True
                    print(f"{CYAN}{event.content}{RESET}", end="", flush=True)

                elif isinstance(event, PlanEvent):
                    # Level 3 Visualization
                    print(
                        f"\n{BLUE}🗺️  [PLAN]: {event.content}{RESET}",
                        end="",
                        flush=True,
                    )

                elif isinstance(event, ContentEvent):
                    if current_mode != "content":
                        if current_mode == "reasoning":
                            print("\n")
                        print(f"\n{GREEN}🤖 [ANSWER]{RESET}\n", end="")
                        current_mode = "content"
                        stats["inference_ok"] = True
                    print(f"{GREEN}{event.content}{RESET}", end="", flush=True)

                elif isinstance(event, DecisionEvent):
                    stats["decision_detected"] = True
                    print(
                        f"\n{MAGENTA}⚡ [DECISION]: {event.to_dict()}{RESET}",
                        end="",
                        flush=True,
                    )

                elif isinstance(event, ToolCallRequestEvent):
                    print(f"\n{YELLOW}🛠  [TOOL DETECTED]: {event.tool_name}{RESET}")

                    # Dynamic Dispatch
                    handler = TOOL_REGISTRY.get(event.tool_name)
                    if handler:
                        # SDK automatically handles output submission
                        if event.execute(handler):
                            stats["tool_executed_count"] += 1
                            seen_tools.add(event.tool_name)
                    else:
                        print(f"{RED} [!] No handler for {event.tool_name}{RESET}")

                elif isinstance(event, StatusEvent) and event.status == "failed":
                    stats["error"] = True
                    stats["detailed_error"] = "Stream Status: Failed"

            print()

        except Exception as e:
            print(f"\n{RED}[!] Stream Error: {e}{RESET}")
            stats["error"] = True
            stats["detailed_error"] = str(e)

        return stats


# ------------------------------------------------------------------
# 4. Main Execution
# ------------------------------------------------------------------
def main(models_to_run):
    if not USER_ID:
        print(f"{RED}CRITICAL: Missing USER_ID in {CONFIG_FILE_NAME} or .env{RESET}")
        return

    print(f"Starting Benchmark for {len(models_to_run)} models...\n")
    ReportManager.update_report([])

    for label, raw_config in models_to_run.items():
        config = (
            raw_config
            if isinstance(raw_config, dict)
            else {"id": raw_config, "provider": "hyperbolic"}
        )
        tester = ModelTester(label, config)

        ok, err = tester.setup()
        if not ok:
            print(f"{RED}Failed setup for {label}: {err}{RESET}")
            ReportManager.update_report(
                [
                    {
                        "name": label,
                        "provider": config["provider"],
                        "id": config["id"],
                        "inference_ok": False,
                        "reasoning_detected": False,
                        "tool_call_ok": False,
                        "batch_ok": False,
                        "call_telemetry": False,
                        "error_msg": err,
                    }
                ]
            )
            continue

        res = tester.run_benchmark()
        ReportManager.update_report([res])
        time.sleep(2)


if __name__ == "__main__":
    main(models_to_run=config.get("models_to_run"))