import json
import os
import re
import time
import urllib.error
import urllib.request
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime

# --- CONFIGURATION ---
TARGET_FILE = "models.py"
MANIFEST_FILE = "model_manifest.json"
REPORT_FILE = "HEALTH_REPORT.md"
CONCURRENCY = 8  # Higher concurrency since threads are lightweight
TIMEOUT_SECONDS = 15


# Load .env manually if dotenv lib is not allowed,
# otherwise we assume env vars are set in the shell or we parse .env crudely here.
def load_env_manual():
    if os.path.exists(".env"):
        with open(".env", "r") as f:
            for line in f:
                if "=" in line and not line.startswith("#"):
                    key, val = line.strip().split("=", 1)
                    os.environ[key] = val.strip('"').strip("'")


load_env_manual()

PROVIDERS = {
    "hyperbolic": {
        "url": "https://api.hyperbolic.xyz/v1/models",
        "chat_url": "https://api.hyperbolic.xyz/v1/chat/completions",
        "api_key": os.getenv("HYPERBOLIC_API_KEY"),
        "exclude": ["StableDiffusion", "FLUX", "TTS", "SDXL", "diffusion"],
    },
    "together-ai": {
        "url": "https://api.together.xyz/v1/models",
        "chat_url": "https://api.together.xyz/v1/chat/completions",
        "api_key": os.getenv("TOGETHER_API_KEY"),
        "exclude": ["image", "video"],
    },
}

TEST_TOOL = {
    "type": "function",
    "function": {
        "name": "get_flight_times",
        "description": "Get flight info",
        "parameters": {
            "type": "object",
            "properties": {"dep": {"type": "string"}, "arr": {"type": "string"}},
            "required": ["dep", "arr"],
        },
    },
}

# --- 1. NETWORK UTILS (Standard Lib) ---


def raw_request(url, method="GET", api_key=None, payload=None):
    """
    Executes a raw HTTP request using urllib (No 'requests' lib needed).
    """
    if not api_key:
        return None, "Missing API Key"

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "User-Agent": "Entities-Gatekeeper/1.0",
    }

    data = json.dumps(payload).encode("utf-8") if payload else None
    req = urllib.request.Request(url, data=data, headers=headers, method=method)

    try:
        with urllib.request.urlopen(req, timeout=TIMEOUT_SECONDS) as response:
            return json.load(response), None
    except urllib.error.HTTPError as e:
        return None, f"HTTP {e.code}: {e.reason}"
    except urllib.error.URLError as e:
        return None, f"Connection Error: {e.reason}"
    except Exception as e:
        return None, str(e)


# --- 2. FETCHERS ---


def fetch_model_list(provider_name):
    print(f"   ⬇️  Fetching {provider_name} list...")
    config = PROVIDERS[provider_name]
    data, err = raw_request(config["url"], "GET", config["api_key"])

    if err:
        print(f"   ❌ {provider_name} failed: {err}")
        return []

    # Unify response formats (Together returns list directly, Hyperbolic wraps in 'data')
    raw_list = data.get("data", data) if isinstance(data, dict) else data

    candidates = []
    for m in raw_list:
        mid = m.get("id")
        mtype = m.get("type", "chat")  # Default to chat if unspecified

        # Filter Logic
        if any(x in mid for x in config["exclude"]):
            continue
        if mtype not in ["chat", "language", "code", None]:
            continue

        candidates.append(
            {
                "id": mid,
                "provider": provider_name,
                "full_id": f"{provider_name}/{mid}",
                "name": mid.split("/")[-1],
            }
        )
    return candidates


# --- 3. VERIFICATION WORKER ---


def verify_model(model_info):
    """
    Performs the 'Dual Detector' check via a live POST request.
    """
    config = PROVIDERS[model_info["provider"]]

    payload = {
        "model": model_info["id"],
        "messages": [
            {"role": "user", "content": "What is the flight time from NYC to LON?"}
        ],
        "tools": [TEST_TOOL],
        "tool_choice": "auto",
        "max_tokens": 300,
        "temperature": 0.1,
    }

    start_t = time.time()
    response_data, error = raw_request(
        config["chat_url"], "POST", config["api_key"], payload
    )
    latency = round(time.time() - start_t, 2)

    result = {
        **model_info,
        "status": "dead",
        "protocol": "none",
        "latency": latency,
        "error": error,
    }

    if error:
        return result

    # Analysis
    try:
        choice = response_data["choices"][0]
        msg = choice.get("message", {})
        content = msg.get("content") or ""
        tool_calls = msg.get("tool_calls")

        result["status"] = "alive"

        # Protocol Detection
        if tool_calls:
            result["protocol"] = "native"
        elif re.search(r"<fc>|<tool_code>|<tool_call>", content):
            result["protocol"] = "hermes"
        else:
            result["protocol"] = "chat-only"

    except (KeyError, IndexError):
        result["error"] = "Malformed Response"

    # Console Feedback
    icon = "✅" if result["status"] == "alive" else "💀"
    print(
        f"{icon} [{model_info['provider'][:4].upper()}] {model_info['name'][:30]:<30} | {result['protocol'].upper():<10} | {latency}s"
    )

    return result


# --- 4. FILE OPERATIONS ---

# --- 4. FILE OPERATIONS (Fixed for Windows) ---


def generate_files(results):
    # 1. Manifest (JSON)
    # Added encoding="utf-8" here
    with open(MANIFEST_FILE, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2)

    # 2. Markdown Report
    alive = [r for r in results if r["status"] == "alive"]

    # Added encoding="utf-8" here to support Emojis on Windows
    with open(REPORT_FILE, "w", encoding="utf-8") as f:
        f.write(f"# Model Health Report\n**Generated:** {datetime.now()}\n\n")
        f.write(f"**Total:** {len(results)} | **Alive:** {len(alive)}\n\n")
        f.write(
            "| Provider | Model | Proto | Latency | Status |\n|---|---|---|---|---|\n"
        )
        for r in sorted(results, key=lambda x: (x["provider"], x["status"])):
            icon = "🟢" if r["status"] == "alive" else "🔴"
            f.write(
                f"| {r['provider']} | {r['name']} | {r['protocol']} | {r['latency']}s | {icon} |\n"
            )

    print("\n📄 Manifest & Report generated.")
    return alive


def inject_python(valid_models):
    if not os.path.exists(TARGET_FILE):
        print(f"⚠️ {TARGET_FILE} not found. Skipping injection.")
        return

    # Helper to build dict string
    def build_dict_str(models):
        grouped = defaultdict(list)
        for m in models:
            # Org grouping
            org = m["id"].split("/")[0] if "/" in m["id"] else "Other"
            grouped[org].append(m)

        lines = ["{"]
        for org in sorted(grouped.keys()):
            lines.append(f"    # --- {org} ---")
            for m in sorted(grouped[org], key=lambda x: x["name"]):
                lines.append(f"    \"{m['name']}\": {{")
                lines.append(f"        \"id\": \"{m['full_id']}\",")
                lines.append(f"        \"provider\": \"{m['provider']}\",")
                lines.append(f"        \"protocol\": \"{m['protocol']}\",")
                lines.append("    },")
            lines.append("")
        lines.append("}")
        return "\n".join(lines)

    # Filter sets
    hyp_str = build_dict_str([m for m in valid_models if m["provider"] == "hyperbolic"])
    tog_str = build_dict_str(
        [m for m in valid_models if m["provider"] == "together-ai"]
    )

    # Added encoding="utf-8" for safe reading
    with open(TARGET_FILE, "r", encoding="utf-8") as f:
        content = f.read()

    # Regex Replacements
    content = re.sub(
        r"(# --- HYPERBOLIC_MODELS_START ---)(.*?)(# --- HYPERBOLIC_MODELS_END ---)",
        f"\\1\nHYPERBOLIC_MODELS = {hyp_str}\n\\3",
        content,
        flags=re.DOTALL,
    )
    content = re.sub(
        r"(# --- TOGETHER_AI_MODELS_START ---)(.*?)(# --- TOGETHER_AI_MODELS_END ---)",
        f"\\1\nTOGETHER_AI_MODELS = {tog_str}\n\\3",
        content,
        flags=re.DOTALL,
    )

    # Added encoding="utf-8" for safe writing
    with open(TARGET_FILE, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"💉 Successfully injected healthy models into {TARGET_FILE}")


# --- MAIN ---


def main():
    print("🚀 Starting Zero-Dependency Gatekeeper...")

    # 1. Fetch Lists
    candidates = []
    candidates.extend(fetch_model_list("hyperbolic"))
    candidates.extend(fetch_model_list("together-ai"))

    print(
        f"📋 Verification started for {len(candidates)} models (Threads={CONCURRENCY})..."
    )

    # 2. Threaded Verification
    results = []
    with ThreadPoolExecutor(max_workers=CONCURRENCY) as executor:
        future_to_model = {executor.submit(verify_model, m): m for m in candidates}
        for future in as_completed(future_to_model):
            results.append(future.result())

    # 3. Output
    alive_models = generate_files(results)
    inject_python(alive_models)


if __name__ == "__main__":
    main()