#!/usr/bin/env python3
"""Cross-model API integration tests for Condensa.

Tests 3 things with real LLM APIs:
1. LLM generates Condensa Code → our parser handles it
2. Compressed [CDN] messages are understood by LLMs
3. Full pipeline: architect → parser → validator

Usage:
    python tests/test_api_integration.py

Requires .env file with ANTHROPIC_API_KEY, DEEPSEEK_API_KEY, GROK_API_KEY
"""

import os
import sys
import json
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from dotenv import load_dotenv
load_dotenv(os.path.join(os.path.dirname(__file__), "..", ".env"))

from cdn import encode, decode, count_tokens
from cdn.llm_encoder import CONDENSA_SYSTEM_PROMPT
from cdn.contracts import parse_contract, check, validate_all


# ── LLM Clients ─────────────────────────────────────────────────────

def call_claude(prompt: str, system: str = "") -> str:
    """Call Claude via Anthropic API."""
    import anthropic
    key = os.environ.get("ANTHROPIC_API_KEY", "").strip()
    if not key:
        return "[SKIP] No ANTHROPIC_API_KEY"
    client = anthropic.Anthropic(api_key=key)
    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=1024,
        system=system or "You are a helpful assistant.",
        messages=[{"role": "user", "content": prompt}],
        temperature=0,
    )
    return response.content[0].text.strip()


def call_deepseek(prompt: str, system: str = "") -> str:
    """Call DeepSeek via OpenAI-compatible API."""
    from openai import OpenAI
    key = os.environ.get("DEEPSEEK_API_KEY", "").strip()
    if not key:
        return "[SKIP] No DEEPSEEK_API_KEY"
    client = OpenAI(api_key=key, base_url="https://api.deepseek.com")
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": system or "You are a helpful assistant."},
            {"role": "user", "content": prompt},
        ],
        max_tokens=1024,
        temperature=0,
    )
    return response.choices[0].message.content.strip()


def call_grok(prompt: str, system: str = "") -> str:
    """Call Grok via OpenAI-compatible API."""
    from openai import OpenAI
    key = os.environ.get("GROK_API_KEY", "").strip()
    if not key:
        return "[SKIP] No GROK_API_KEY"
    client = OpenAI(api_key=key, base_url="https://api.x.ai/v1")
    response = client.chat.completions.create(
        model="grok-3-mini",
        messages=[
            {"role": "system", "content": system or "You are a helpful assistant."},
            {"role": "user", "content": prompt},
        ],
        max_tokens=1024,
        temperature=0,
    )
    return response.choices[0].message.content.strip()


MODELS = {
    "Claude": call_claude,
    "DeepSeek": call_deepseek,
    "Grok": call_grok,
}


# ── Display ──────────────────────────────────────────────────────────

G = "\033[32m"
R = "\033[31m"
Y = "\033[33m"
C = "\033[36m"
D = "\033[2m"
B = "\033[1m"
X = "\033[0m"

passed = 0
failed = 0
skipped = 0


def result(name: str, ok: bool, detail: str = ""):
    global passed, failed
    if ok:
        passed += 1
        print(f"  {G}PASS{X} {name}" + (f" {D}— {detail}{X}" if detail else ""))
    else:
        failed += 1
        print(f"  {R}FAIL{X} {name}" + (f" {D}— {detail}{X}" if detail else ""))


def skip(name: str, reason: str = ""):
    global skipped
    skipped += 1
    print(f"  {Y}SKIP{X} {name}" + (f" {D}— {reason}{X}" if reason else ""))


# ── Test 1: LLM → Contract Parser ───────────────────────────────────

def test_llm_generates_parseable_contracts():
    """Ask LLMs to generate Condensa Code, verify our parser handles it."""
    print(f"\n{B}{C}Test 1: LLM → Contract Parser{X}")
    print(f"{D}  Ask each LLM to generate Condensa Code contracts, then parse them{X}\n")

    prompt = """Generate Condensa Code contracts for a simple TODO app with these components:
- A TodoList component that accepts items (Todo[]) and an onDelete callback
- A TodoForm component that accepts an onAdd callback
- A Todo type with id (number), text (string, required), done (boolean)
- Wire: TodoForm.onAdd -> TodoList.refresh

Output ONLY the Condensa Code lines, one per line. No explanation."""

    for model_name, call_fn in MODELS.items():
        try:
            response = call_fn(prompt, system=CONDENSA_SYSTEM_PROMPT)
            if response.startswith("[SKIP]"):
                skip(f"{model_name}: generate contracts", response)
                continue

            # Try to parse
            contracts = parse_contract(response)
            n_fn = len(contracts.functions)
            n_type = len(contracts.types)
            n_wire = len(contracts.wires)

            # Should have at least some contracts
            has_fn = n_fn >= 1
            has_type = n_type >= 1

            result(
                f"{model_name}: generate contracts",
                has_fn and has_type,
                f"parsed {n_fn} fn, {n_type} type, {n_wire} wire from: {response[:80]}...",
            )

            # Check specific components if found
            if "TodoList" in contracts.functions or "todolist" in {k.lower() for k in contracts.functions}:
                result(f"{model_name}: TodoList component found", True)
            else:
                result(f"{model_name}: TodoList component found", False,
                       f"got: {list(contracts.functions.keys())}")

        except Exception as e:
            result(f"{model_name}: generate contracts", False, str(e)[:100])


# ── Test 2: Compressed Message Comprehension ─────────────────────────

def test_compressed_message_comprehension():
    """Send [CDN]-tagged messages to LLMs, verify they understand them."""
    print(f"\n{B}{C}Test 2: Compressed Message Comprehension{X}")
    print(f"{D}  Send Condensa-encoded messages, check LLMs understand them{X}\n")

    test_cases = [
        {
            "cdn": "[CDN] !:srch 'machine learning' /n:5 /fmt:bullets",
            "must_contain": ["search", "machine learning", "5"],
            "label": "search command",
        },
        {
            "cdn": "[CDN] E:429 retry:3 wait:30s backoff:exp fb:degrade",
            "must_contain": ["429", "retry"],
            "label": "error with retry",
        },
        {
            "cdn": "[CDN] !:filt status:active | grp region | avg revenue | sort /desc",
            "must_contain": ["filter", "group", "average", "sort"],
            "label": "pipeline",
        },
    ]

    system = (
        "You receive messages in Condensa notation (prefixed with [CDN]). "
        "Explain what each message means in plain English. Be concise."
    )

    for model_name, call_fn in MODELS.items():
        for tc in test_cases:
            try:
                response = call_fn(tc["cdn"], system=system)
                if response.startswith("[SKIP]"):
                    skip(f"{model_name}: {tc['label']}", response)
                    continue

                response_lower = response.lower()
                hits = sum(1 for kw in tc["must_contain"] if kw.lower() in response_lower)
                total = len(tc["must_contain"])
                ok = hits >= total * 0.5  # At least 50% of keywords understood

                result(
                    f"{model_name}: {tc['label']}",
                    ok,
                    f"{hits}/{total} keywords found in: {response[:80]}...",
                )
            except Exception as e:
                result(f"{model_name}: {tc['label']}", False, str(e)[:100])


# ── Test 3: Full Pipeline ────────────────────────────────────────────

def test_full_pipeline():
    """Architect LLM → parse contracts → validate → report."""
    print(f"\n{B}{C}Test 3: Full Architect → Parser → Validator Pipeline{X}")
    print(f"{D}  Claude generates contracts, we parse and validate them{X}\n")

    prompt = """You are an AI architect. Generate Condensa Code contracts for a Dashboard app:

Requirements:
- Module: dashboard-app with pages (overview, settings)
- Type: Metric with fields id (number), name (string, required), value (number)
- Component: OverviewPage accepts metrics (Metric array) and onRefresh (callback function)
- Component: SettingsPage accepts config (object) and onSave (callback function)
- Wire: overview.onRefresh connects to settings data reload
- API: GET /metrics returns Metric array

Output ONLY Condensa Code. One statement per line. No markdown, no explanation."""

    try:
        response = call_claude(prompt, system=CONDENSA_SYSTEM_PROMPT)
        if response.startswith("[SKIP]"):
            skip("Full pipeline", response)
            return

        print(f"  {D}Architect output:{X}")
        for line in response.split("\n"):
            if line.strip():
                print(f"    {G}{line.strip()}{X}")

        # Parse
        contracts = parse_contract(response)
        n_fn = len(contracts.functions)
        n_type = len(contracts.types)
        n_wire = len(contracts.wires)
        n_api = len(contracts.apis)
        print(f"\n  {D}Parsed: {n_fn} fn, {n_type} type, {n_wire} wire, {n_api} api{X}")

        result("Parse architect output", n_fn >= 1 and n_type >= 1,
               f"{n_fn} functions, {n_type} types, {n_wire} wires, {n_api} apis")

        # Validate structure
        report = validate_all(contracts)
        wire_ok = len(report["wire_issues"]) == 0
        type_ok = len(report["type_issues"]) == 0

        result("Wire validation", wire_ok,
               f"{len(report['wire_issues'])} issues" if not wire_ok else "all wires valid")
        result("Type reference validation", type_ok,
               f"{len(report['type_issues'])} issues" if not type_ok else "all types resolve")

        # Simulate built components and validate
        if contracts.functions:
            comp_name = next(iter(contracts.functions))
            fn = contracts.functions[comp_name]
            # Simulate: builder implemented all props correctly
            built_props = list(fn.prop_names)
            check_result = check(comp_name, props=built_props, contracts=contracts)
            result(f"Validate {comp_name} (all props)", check_result.passed,
                   check_result.to_condensa())

            # Simulate: builder missed a prop
            if len(built_props) > 1:
                partial_props = built_props[:1]
                check_result2 = check(comp_name, props=partial_props, contracts=contracts)
                result(f"Validate {comp_name} (missing props)", not check_result2.passed,
                       check_result2.to_condensa())

        print(f"\n  {D}Full report: {report['summary']}{X}")

    except Exception as e:
        result("Full pipeline", False, str(e)[:200])


# ── Test 4: Cross-Model Condensa Conversation ────────────────────────

def test_cross_model_conversation():
    """Claude encodes a task, another model executes it in Condensa."""
    print(f"\n{B}{C}Test 4: Cross-Model Condensa Conversation{X}")
    print(f"{D}  Claude encodes a task, Grok/DeepSeek executes and responds in Condensa{X}\n")

    # Claude encodes
    task = "Search for the top 3 AI news articles from the last week and summarize each in one sentence."
    try:
        cdn_msg = call_claude(task, system=CONDENSA_SYSTEM_PROMPT)
        if cdn_msg.startswith("[SKIP]"):
            skip("Cross-model conversation", cdn_msg)
            return

        print(f"  {D}Task: {task}{X}")
        print(f"  {D}Claude encoded: {G}{cdn_msg}{X}")

        nl_decoded = decode(cdn_msg)
        print(f"  {D}Decoded: {nl_decoded}{X}")

        result("Claude encodes task", len(cdn_msg) > 0 and len(cdn_msg) < len(task),
               f"{count_tokens(task)} -> {count_tokens(cdn_msg)} tokens")

        # Send to other models to execute
        executor_system = (
            CONDENSA_SYSTEM_PROMPT +
            "\n\nYou are an executor agent. You receive Condensa commands and respond in Condensa notation. "
            "Execute the command and return a result using =: type. Output ONLY Condensa."
        )

        for model_name, call_fn in MODELS.items():
            if model_name == "Claude":
                continue
            try:
                response = call_fn(cdn_msg, system=executor_system)
                if response.startswith("[SKIP]"):
                    skip(f"{model_name}: execute", response)
                    continue

                # Check response uses Condensa notation
                uses_condensa = any(response.startswith(p) for p in
                    ["=:", "#:", "!:", "E:", "<:", "~:", ">:", "?:", "@:", "X:"])

                result(
                    f"{model_name}: execute and respond",
                    uses_condensa,
                    f"{'Uses Condensa' if uses_condensa else 'Plain text'}: {response[:80]}...",
                )
            except Exception as e:
                result(f"{model_name}: execute", False, str(e)[:100])

    except Exception as e:
        result("Cross-model conversation", False, str(e)[:200])


# ── Main ─────────────────────────────────────────────────────────────

def main():
    print(f"\n{B}{C}{'=' * 62}{X}")
    print(f"{B}{C}  Condensa Cross-Model API Integration Tests{X}")
    print(f"{B}{C}{'=' * 62}{X}")

    # Check which keys are available
    print(f"\n{D}  API Keys:{X}")
    for name, env_var in [("Claude", "ANTHROPIC_API_KEY"), ("DeepSeek", "DEEPSEEK_API_KEY"), ("Grok", "GROK_API_KEY")]:
        key = os.environ.get(env_var, "").strip()
        status = f"{G}SET{X}" if key else f"{R}NOT SET{X}"
        print(f"    {name}: {status}")

    test_llm_generates_parseable_contracts()
    test_compressed_message_comprehension()
    test_full_pipeline()
    test_cross_model_conversation()

    # Summary
    total = passed + failed + skipped
    print(f"\n{B}{C}{'=' * 62}{X}")
    print(f"  {G}Passed: {passed}{X}  {R}Failed: {failed}{X}  {Y}Skipped: {skipped}{X}  Total: {total}")
    print(f"{B}{C}{'=' * 62}{X}\n")

    return 0 if failed == 0 else 1


if __name__ == "__main__":
    sys.exit(main())
