"""
Test: Do LLMs generate correct async tool code on the first try?

Sends tasks to GPT-4.1-mini with ez-ptc async tool prompts, then runs the
generated code through toolkit.execute_sync() and reports success/failure.

Tests 3 tasks of varying complexity:
  1. Simple: single async tool call
  2. Medium: two sequential async tool calls
  3. Harder: parallel async tool calls with comparison
"""

import asyncio
import json
import os
import time
from typing import TypedDict

from dotenv import load_dotenv
from openai import OpenAI

from ez_ptc import Toolkit, ez_tool

load_dotenv()


# ── Async mock tools ──────────────────────────────────────────────────────


class WeatherResult(TypedDict):
    city: str
    temp_celsius: float
    condition: str
    humidity_pct: int


class SearchResult(TypedDict):
    item_id: str
    title: str
    relevance_score: float


class ItemDetails(TypedDict):
    item_id: str
    title: str
    description: str
    tags: list[str]
    rating: float


MOCK_WEATHER = {
    "tokyo": {"city": "Tokyo", "temp_celsius": 18.5, "condition": "Cloudy", "humidity_pct": 72},
    "new york": {"city": "New York", "temp_celsius": 25.0, "condition": "Sunny", "humidity_pct": 55},
    "london": {"city": "London", "temp_celsius": 14.0, "condition": "Rainy", "humidity_pct": 88},
    "paris": {"city": "Paris", "temp_celsius": 20.0, "condition": "Partly Cloudy", "humidity_pct": 65},
}

MOCK_SEARCH = {
    "python": [
        {"item_id": "py_001", "title": "Python Crash Course", "relevance_score": 0.95},
        {"item_id": "py_002", "title": "Fluent Python", "relevance_score": 0.90},
        {"item_id": "py_003", "title": "Python Data Science Handbook", "relevance_score": 0.85},
    ],
    "rust": [
        {"item_id": "rs_001", "title": "The Rust Programming Language", "relevance_score": 0.98},
        {"item_id": "rs_002", "title": "Rust in Action", "relevance_score": 0.88},
    ],
}

MOCK_DETAILS = {
    "py_001": {
        "item_id": "py_001",
        "title": "Python Crash Course",
        "description": "A hands-on, project-based introduction to programming",
        "tags": ["beginner", "python", "projects"],
        "rating": 4.7,
    },
    "py_002": {
        "item_id": "py_002",
        "title": "Fluent Python",
        "description": "Clear, concise, and effective programming in Python",
        "tags": ["advanced", "python", "idioms"],
        "rating": 4.8,
    },
    "py_003": {
        "item_id": "py_003",
        "title": "Python Data Science Handbook",
        "description": "Essential tools for working with data in Python",
        "tags": ["data-science", "python", "numpy", "pandas"],
        "rating": 4.5,
    },
    "rs_001": {
        "item_id": "rs_001",
        "title": "The Rust Programming Language",
        "description": "The official Rust language book",
        "tags": ["rust", "systems", "official"],
        "rating": 4.9,
    },
    "rs_002": {
        "item_id": "rs_002",
        "title": "Rust in Action",
        "description": "Systems programming concepts and techniques with Rust",
        "tags": ["rust", "systems", "intermediate"],
        "rating": 4.6,
    },
}


@ez_tool
async def get_weather(city: str) -> WeatherResult:
    """Get current weather for a city. Returns temperature, condition, and humidity."""
    await asyncio.sleep(0)  # simulate async I/O
    key = city.lower().strip()
    if key not in MOCK_WEATHER:
        return {"city": city, "temp_celsius": 0.0, "condition": "Unknown", "humidity_pct": 0}
    return MOCK_WEATHER[key]


@ez_tool
async def search_items(query: str) -> list[SearchResult]:
    """Search for items by keyword. Returns a list of matching items with relevance scores."""
    await asyncio.sleep(0)  # simulate async I/O
    key = query.lower().strip()
    if key not in MOCK_SEARCH:
        return []
    return MOCK_SEARCH[key]


@ez_tool
async def get_item_details(item_id: str) -> ItemDetails:
    """Get detailed information about a specific item by its item_id."""
    await asyncio.sleep(0)  # simulate async I/O
    if item_id not in MOCK_DETAILS:
        raise ValueError(f"Item '{item_id}' not found")
    return MOCK_DETAILS[item_id]


TOOLS = [get_weather, search_items, get_item_details]


# ── Tasks ─────────────────────────────────────────────────────────────────

TASKS = [
    {
        "name": "Task 1 (simple): Single async tool call",
        "user_msg": "Get the current weather in Tokyo.",
    },
    {
        "name": "Task 2 (medium): Two sequential async tool calls",
        "user_msg": "Search for items about 'python', then get the full details of the first result.",
    },
    {
        "name": "Task 3 (harder): Parallel async tool calls + comparison",
        "user_msg": (
            "Get the weather in Tokyo and New York at the same time (in parallel), "
            "then print which city is warmer and by how many degrees."
        ),
    },
]


# ── Runner ────────────────────────────────────────────────────────────────


def run_test(client: OpenAI, model: str, toolkit: Toolkit, task: dict) -> dict:
    """Send a task to the LLM via tool mode, execute the code, report results."""
    openai_tool_schema = toolkit.tool_schema(format="openai")
    system_prompt = (
        "You are a helpful assistant.\n\n" + toolkit.tool_prompt()
    )

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": task["user_msg"]},
    ]

    # Single LLM call — we want first-try success
    start = time.perf_counter()
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        tools=[openai_tool_schema],
        temperature=0.0,
    )
    api_time = time.perf_counter() - start

    choice = response.choices[0]
    message = choice.message

    # Extract code from tool call
    if not message.tool_calls:
        return {
            "success": False,
            "code": None,
            "error": "LLM did not make a tool call — responded with plain text instead.",
            "llm_text": message.content or "",
            "output": "",
            "api_time": api_time,
        }

    args = json.loads(message.tool_calls[0].function.arguments)
    code = args.get("code", "")

    # Execute the code
    result = toolkit.execute_sync(code)

    return {
        "success": result.success,
        "code": code,
        "error": result.error if not result.success else None,
        "error_output": result.error_output if not result.success else None,
        "output": result.output,
        "tool_calls": [tc.to_dict() if hasattr(tc, "to_dict") else str(tc) for tc in result.tool_calls],
        "api_time": api_time,
    }


def diagnose_failure(code: str, error: str, error_output: str | None) -> list[str]:
    """Identify common LLM failure modes from the code and error."""
    issues = []
    if code is None:
        issues.append("LLM did not generate any code")
        return issues

    if "asyncio.run(" in code:
        issues.append("Used asyncio.run() inside already-async context")
    if "import " in code and any(
        f"import {t}" in code for t in ["get_weather", "search_items", "get_item_details"]
    ):
        issues.append("Tried to import pre-loaded tool functions")
    if "await" not in code and any(
        t in code for t in ["get_weather(", "search_items(", "get_item_details("]
    ):
        issues.append("Called async tool without await")

    if error_output:
        if "KeyError" in error_output:
            issues.append("KeyError — LLM guessed wrong dictionary key")
        if "AttributeError" in error_output:
            issues.append("AttributeError — LLM used dot-access on dict or wrong attribute")
        if "NameError" in error_output:
            issues.append("NameError — LLM referenced undefined variable or tried to use unimported module")
        if "TypeError" in error_output:
            issues.append("TypeError — wrong argument types or calling convention")
        if "Validation failed" in error_output or "Validation errors" in error_output:
            issues.append("Code failed AST validation (blocked import, unsafe operation, etc.)")

    if not issues:
        issues.append(f"Unknown failure mode — error: {error}")

    return issues


def main():
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        print("ERROR: OPENAI_API_KEY not found in environment. Check .env file.")
        return

    model = "gpt-4.1-mini"
    client = OpenAI(api_key=api_key)

    # Test both chaining=True and chaining=False
    for chaining in [True, False]:
        chaining_label = "chaining=True" if chaining else "chaining=False"
        toolkit = Toolkit(TOOLS, assist_tool_chaining=chaining)

        print(f"\n{'#' * 80}")
        print(f"# ASYNC TOOL TEST — {chaining_label}")
        print(f"# Model: {model}")
        print(f"{'#' * 80}")

        # Print the prompt for reference
        print(f"\n--- Tool Prompt ({chaining_label}) ---")
        prompt = toolkit.tool_prompt()
        for line in prompt.split("\n"):
            print(f"  {line}")
        print()

        results = []
        for task in TASKS:
            print(f"\n{'=' * 70}")
            print(f"  {task['name']} ({chaining_label})")
            print(f"  User: {task['user_msg']}")
            print(f"{'=' * 70}")

            result = run_test(client, model, toolkit, task)
            results.append(result)

            print(f"\n  --- LLM-Generated Code ---")
            if result["code"]:
                for line in result["code"].split("\n"):
                    print(f"    {line}")
            else:
                print(f"    (no code generated)")

            print(f"\n  --- Execution Result ---")
            status = "SUCCESS" if result["success"] else "FAILED"
            print(f"    Status: {status}")
            print(f"    API Time: {result['api_time']:.2f}s")

            if result["success"]:
                print(f"    Output:")
                for line in result["output"].split("\n"):
                    print(f"      {line}")
            else:
                print(f"    Error: {result.get('error', 'N/A')}")
                if result.get("error_output"):
                    print(f"    Error Output:")
                    for line in result["error_output"].split("\n"):
                        print(f"      {line}")
                # Diagnose
                issues = diagnose_failure(
                    result["code"],
                    result.get("error", ""),
                    result.get("error_output"),
                )
                print(f"    Diagnosis:")
                for issue in issues:
                    print(f"      - {issue}")

        # Summary
        successes = sum(1 for r in results if r["success"])
        total = len(results)
        print(f"\n{'=' * 70}")
        print(f"  SUMMARY ({chaining_label}): {successes}/{total} tasks succeeded on first try")
        print(f"{'=' * 70}")

        if successes < total:
            print(f"\n  Failed tasks:")
            for i, (task, result) in enumerate(zip(TASKS, results)):
                if not result["success"]:
                    issues = diagnose_failure(
                        result["code"],
                        result.get("error", ""),
                        result.get("error_output"),
                    )
                    print(f"    - {task['name']}: {', '.join(issues)}")


if __name__ == "__main__":
    main()
