"""
Three-Way Comparison: Traditional vs ez-ptc (no chaining) vs ez-ptc (with chaining)

Compares:
1. Traditional tool calling (Pydantic AI with individual tools)
2. ez-ptc without tool chaining (assist_tool_chaining=False)
3. ez-ptc with tool chaining (assist_tool_chaining=True)
"""

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

from dotenv import load_dotenv
from ez_ptc import Toolkit, ez_tool
from pydantic_ai import Agent, Tool
from pydantic_ai.models.openai import OpenAIResponsesModelSettings

load_dotenv()


# Type definitions
class CarSearchResult(TypedDict):
    car_id: str
    name: str
    price: int
    body_type: str


class CarSpecs(TypedDict):
    car_id: str
    kmpl: float
    fuel_type: str
    engine_cc: int
    transmission: str


class AvailabilityResult(TypedDict):
    car_id: str
    available: bool
    dealer_count: int


# Mock data — 15 SUVs
MOCK_CARS = [
    {"car_id": "suv_001", "name": "Hyundai Creta", "price": 1100000, "body_type": "SUV"},
    {"car_id": "suv_002", "name": "Kia Seltos", "price": 1200000, "body_type": "SUV"},
    {"car_id": "suv_003", "name": "MG Hector", "price": 1450000, "body_type": "SUV"},
    {"car_id": "suv_004", "name": "Tata Nexon", "price": 900000, "body_type": "SUV"},
    {"car_id": "suv_005", "name": "Maruti Brezza", "price": 850000, "body_type": "SUV"},
    {"car_id": "suv_006", "name": "Mahindra XUV300", "price": 950000, "body_type": "SUV"},
    {"car_id": "suv_007", "name": "Toyota Urban Cruiser", "price": 1000000, "body_type": "SUV"},
    {"car_id": "suv_008", "name": "Volkswagen Taigun", "price": 1150000, "body_type": "SUV"},
    {"car_id": "suv_009", "name": "Skoda Kushaq", "price": 1180000, "body_type": "SUV"},
    {"car_id": "suv_010", "name": "Honda WR-V", "price": 980000, "body_type": "SUV"},
    {"car_id": "suv_011", "name": "Renault Kiger", "price": 650000, "body_type": "SUV"},
    {"car_id": "suv_012", "name": "Nissan Magnite", "price": 620000, "body_type": "SUV"},
    {"car_id": "suv_013", "name": "Mahindra XUV700", "price": 1400000, "body_type": "SUV"},
    {"car_id": "suv_014", "name": "Tata Harrier", "price": 1500000, "body_type": "SUV"},
    {"car_id": "suv_015", "name": "Hyundai Venue", "price": 780000, "body_type": "SUV"},
]

MOCK_SPECS = {
    "suv_001": {"car_id": "suv_001", "kmpl": 16.8, "fuel_type": "Petrol", "engine_cc": 1497, "transmission": "Manual"},
    "suv_002": {"car_id": "suv_002", "kmpl": 16.5, "fuel_type": "Petrol", "engine_cc": 1497, "transmission": "Automatic"},
    "suv_003": {"car_id": "suv_003", "kmpl": 13.5, "fuel_type": "Petrol", "engine_cc": 1956, "transmission": "Automatic"},
    "suv_004": {"car_id": "suv_004", "kmpl": 17.4, "fuel_type": "Diesel", "engine_cc": 1497, "transmission": "Manual"},
    "suv_005": {"car_id": "suv_005", "kmpl": 17.0, "fuel_type": "Petrol", "engine_cc": 1462, "transmission": "Manual"},
    "suv_006": {"car_id": "suv_006", "kmpl": 17.0, "fuel_type": "Petrol", "engine_cc": 1197, "transmission": "Manual"},
    "suv_007": {"car_id": "suv_007", "kmpl": 17.0, "fuel_type": "Petrol", "engine_cc": 1462, "transmission": "Automatic"},
    "suv_008": {"car_id": "suv_008", "kmpl": 14.1, "fuel_type": "Petrol", "engine_cc": 999, "transmission": "Automatic"},
    "suv_009": {"car_id": "suv_009", "kmpl": 14.5, "fuel_type": "Petrol", "engine_cc": 999, "transmission": "Manual"},
    "suv_010": {"car_id": "suv_010", "kmpl": 16.5, "fuel_type": "Petrol", "engine_cc": 1199, "transmission": "Manual"},
    "suv_011": {"car_id": "suv_011", "kmpl": 20.5, "fuel_type": "Petrol", "engine_cc": 999, "transmission": "Manual"},
    "suv_012": {"car_id": "suv_012", "kmpl": 20.0, "fuel_type": "Petrol", "engine_cc": 999, "transmission": "Manual"},
    "suv_013": {"car_id": "suv_013", "kmpl": 13.0, "fuel_type": "Petrol", "engine_cc": 1997, "transmission": "Automatic"},
    "suv_014": {"car_id": "suv_014", "kmpl": 14.6, "fuel_type": "Diesel", "engine_cc": 1956, "transmission": "Automatic"},
    "suv_015": {"car_id": "suv_015", "kmpl": 17.5, "fuel_type": "Petrol", "engine_cc": 1197, "transmission": "Manual"},
}

MOCK_AVAILABILITY = {
    "suv_001": {"available": True, "dealer_count": 5},
    "suv_002": {"available": True, "dealer_count": 4},
    "suv_003": {"available": False, "dealer_count": 0},
    "suv_004": {"available": True, "dealer_count": 8},
    "suv_005": {"available": True, "dealer_count": 6},
    "suv_006": {"available": True, "dealer_count": 3},
    "suv_007": {"available": False, "dealer_count": 0},
    "suv_008": {"available": True, "dealer_count": 2},
    "suv_009": {"available": True, "dealer_count": 3},
    "suv_010": {"available": False, "dealer_count": 0},
    "suv_011": {"available": True, "dealer_count": 7},
    "suv_012": {"available": True, "dealer_count": 5},
    "suv_013": {"available": True, "dealer_count": 4},
    "suv_014": {"available": False, "dealer_count": 0},
    "suv_015": {"available": True, "dealer_count": 6},
}


# Tools
@ez_tool
def search_cars(body_type: str) -> list[CarSearchResult]:
    """Search cars by body type."""
    return [
        car for car in MOCK_CARS
        if car["body_type"].lower() == body_type.lower()
    ]


@ez_tool
def get_specs(car_id: str) -> CarSpecs:
    """Get detailed specifications for a specific car."""
    if car_id not in MOCK_SPECS:
        raise ValueError(f"Car ID {car_id} not found")
    return MOCK_SPECS[car_id]


@ez_tool
def check_availability(car_id: str, city: str) -> AvailabilityResult:
    """Check dealer availability for a car in a specific city."""
    if car_id not in MOCK_AVAILABILITY:
        return {"car_id": car_id, "available": False, "dealer_count": 0}
    result = MOCK_AVAILABILITY[car_id]
    return {"car_id": car_id, **result}


QUERY = "SUVs with mileage >15 kmpl, available in Delhi with specs and dealer availability."


def extract_code_blocks(messages_json):
    """Extract LLM-generated code from messages.

    For traditional tool calling: formats each tool call as a function call.
    For ez-ptc: extracts the `code` argument from execute_tools calls.
    Returns a list of code strings (one per turn that contained tool calls).
    """
    blocks = []
    for msg in messages_json:
        if msg.get("kind") != "response":
            continue
        parts = msg.get("parts", [])
        tool_calls_in_turn = []
        for part in parts:
            if part.get("part_kind") != "tool-call":
                continue
            tool_name = part.get("tool_name", "")
            args_raw = part.get("args", "")

            if tool_name == "execute_tools":
                # ez-ptc meta-tool — extract the code argument
                try:
                    args_obj = json.loads(args_raw) if isinstance(args_raw, str) else args_raw
                    code = args_obj.get("code", args_raw)
                except (json.JSONDecodeError, TypeError, AttributeError):
                    code = args_raw
                tool_calls_in_turn.append(code)
            else:
                # Traditional tool call — format as function invocation
                try:
                    args_obj = json.loads(args_raw) if isinstance(args_raw, str) else args_raw
                    if isinstance(args_obj, dict):
                        params = ", ".join(f"{k}={json.dumps(v)}" for k, v in args_obj.items())
                    else:
                        params = str(args_obj)
                except (json.JSONDecodeError, TypeError):
                    params = str(args_raw)
                tool_calls_in_turn.append(f"{tool_name}({params})")

        if tool_calls_in_turn:
            blocks.append("\n".join(tool_calls_in_turn))

    return blocks


def extract_tool_io(messages_json):
    """Extract tool call inputs and outputs from Pydantic AI messages."""
    tool_io = []
    # Build a map of tool_call_id -> args for matching returns
    pending_calls = {}

    for msg in messages_json:
        parts = msg.get("parts", [])
        for part in parts:
            kind = part.get("part_kind", "")

            if kind == "tool-call":
                call_id = part.get("tool_call_id", "")
                entry = {
                    "tool": part.get("tool_name", ""),
                    "input": part.get("args", ""),
                    "output": None,
                }
                pending_calls[call_id] = entry
                tool_io.append(entry)

            elif kind == "tool-return":
                call_id = part.get("tool_call_id", "")
                content = part.get("content", "")
                if call_id in pending_calls:
                    pending_calls[call_id]["output"] = content

    return tool_io


def print_tool_io(tool_io):
    """Pretty-print tool call inputs and outputs."""
    if not tool_io:
        print("  (no tool calls)")
        return
    for i, call in enumerate(tool_io, 1):
        args_str = call["input"]
        if isinstance(args_str, str):
            try:
                args_str = json.dumps(json.loads(args_str), indent=2)
            except (json.JSONDecodeError, TypeError):
                pass
        output_str = call["output"]
        if isinstance(output_str, (dict, list)):
            output_str = json.dumps(output_str, indent=2)
        elif isinstance(output_str, str) and len(output_str) > 200:
            output_str = output_str[:200] + "..."

        print(f"  [{i}] {call['tool']}")
        print(f"      Input:  {args_str}")
        print(f"      Output: {output_str}")


def extract_metrics(result):
    """Extract metrics from a Pydantic AI result."""
    messages_json = json.loads(result.all_messages_json())

    model_responses = 0
    tool_calls = 0
    total_prompt_tokens = 0
    total_completion_tokens = 0

    for msg in messages_json:
        msg_kind = msg.get('kind', '')

        if msg_kind in ['response', 'model-text-response']:
            model_responses += 1

        if msg_kind == 'tool-return':
            tool_calls += 1

        if 'usage' in msg:
            usage = msg['usage']
            if isinstance(usage, dict):
                total_prompt_tokens += usage.get('prompt_tokens', 0)
                total_completion_tokens += usage.get('completion_tokens', 0)
                total_prompt_tokens += usage.get('input_tokens', 0)
                total_completion_tokens += usage.get('output_tokens', 0)

    return {
        "turns": model_responses,
        "tool_calls": tool_calls,
        "prompt_tokens": total_prompt_tokens,
        "completion_tokens": total_completion_tokens,
        "total_tokens": total_prompt_tokens + total_completion_tokens,
        "messages_json": messages_json,
        "tool_io": extract_tool_io(messages_json),
        "code_blocks": extract_code_blocks(messages_json),
    }


async def run_scenario_1_traditional():
    """Scenario 1: Traditional tool calling."""
    print("\n" + "="*80)
    print("SCENARIO 1: Traditional Tool Calling")
    print("="*80 + "\n")

    agent = Agent(
        "openai:gpt-4.1-mini",
        system_prompt="""You are a helpful car search assistant.
Use the available tools to answer the user's query.""",
        tools=[search_cars, get_specs, check_availability],
        model_settings=OpenAIResponsesModelSettings(
            parallel_tool_calls=True
        ),
        retries=3
    )

    start_time = time.time()
    result = await agent.run(QUERY)
    end_time = time.time()

    metrics = extract_metrics(result)
    metrics["time"] = end_time - start_time
    metrics["output"] = result.output if hasattr(result, 'output') else str(result)

    print("--- Generated Code (per turn) ---")
    for i, block in enumerate(metrics["code_blocks"], 1):
        print(f"  Turn {i}:")
        for line in block.split("\n"):
            print(f"    {line}")
    print()
    print("--- Tool Calls ---")
    print_tool_io(metrics["tool_io"])
    print()
    print(f"Time: {metrics['time']:.2f}s")
    print(f"LLM Turns: {metrics['turns']}")
    print(f"Tool Calls: {metrics['tool_calls']}")
    print(f"Tokens: {metrics['total_tokens']:,} (input: {metrics['prompt_tokens']:,}, output: {metrics['completion_tokens']:,})")

    return metrics


async def run_scenario_2_ezptc_no_chaining():
    """Scenario 2: ez-ptc without tool chaining."""
    print("\n" + "="*80)
    print("SCENARIO 2: ez-ptc WITHOUT Tool Chaining (assist_tool_chaining=False)")
    print("="*80 + "\n")

    toolkit = Toolkit(
        [search_cars, get_specs, check_availability],
        assist_tool_chaining=False,
    )

    execute_fn = toolkit.as_tool()
    pydantic_tool = Tool(execute_fn, takes_ctx=False)

    agent = Agent(
        "openai:gpt-4.1-mini",
        system_prompt="You are a helpful car search assistant.",
        tools=[pydantic_tool],
        model_settings=OpenAIResponsesModelSettings(
            parallel_tool_calls=False
        )
    )

    start_time = time.time()
    result = await agent.run(QUERY)
    end_time = time.time()

    metrics = extract_metrics(result)
    metrics["time"] = end_time - start_time
    metrics["output"] = result.output if hasattr(result, 'output') else str(result)

    print("--- Generated Code (per turn) ---")
    for i, block in enumerate(metrics["code_blocks"], 1):
        print(f"  Turn {i}:")
        for line in block.split("\n"):
            print(f"    {line}")
    print()
    print("--- Tool Calls ---")
    print_tool_io(metrics["tool_io"])
    print()
    print(f"Time: {metrics['time']:.2f}s")
    print(f"LLM Turns: {metrics['turns']}")
    print(f"Tool Calls: {metrics['tool_calls']}")
    print(f"Tokens: {metrics['total_tokens']:,} (input: {metrics['prompt_tokens']:,}, output: {metrics['completion_tokens']:,})")

    return metrics


async def run_scenario_3_ezptc_with_chaining():
    """Scenario 3: ez-ptc with tool chaining."""
    print("\n" + "="*80)
    print("SCENARIO 3: ez-ptc WITH Tool Chaining (assist_tool_chaining=True)")
    print("="*80 + "\n")

    toolkit = Toolkit(
        [search_cars, get_specs, check_availability],
        assist_tool_chaining=True
    )

    execute_fn = toolkit.as_tool()
    pydantic_tool = Tool(execute_fn, takes_ctx=False)

    agent = Agent(
        "openai:gpt-4.1-mini",
        system_prompt="You are a helpful car search assistant.",
        tools=[pydantic_tool],
        model_settings=OpenAIResponsesModelSettings(
            parallel_tool_calls=False
        )
    )

    start_time = time.time()
    result = await agent.run(QUERY)
    end_time = time.time()

    metrics = extract_metrics(result)
    metrics["time"] = end_time - start_time
    metrics["output"] = result.output if hasattr(result, 'output') else str(result)

    print("--- Generated Code (per turn) ---")
    for i, block in enumerate(metrics["code_blocks"], 1):
        print(f"  Turn {i}:")
        for line in block.split("\n"):
            print(f"    {line}")
    print()
    print("--- Tool Calls ---")
    print_tool_io(metrics["tool_io"])
    print()
    print(f"Time: {metrics['time']:.2f}s")
    print(f"LLM Turns: {metrics['turns']}")
    print(f"Tool Calls: {metrics['tool_calls']}")
    print(f"Tokens: {metrics['total_tokens']:,} (input: {metrics['prompt_tokens']:,}, output: {metrics['completion_tokens']:,})")

    return metrics


async def main():
    """Run all three scenarios and generate comparison report."""

    if not os.getenv("OPENAI_API_KEY"):
        print("Error: OPENAI_API_KEY not found")
        return

    print("\n" + "="*80)
    print("THREE-WAY COMPARISON")
    print("="*80)

    s1 = await run_scenario_1_traditional()
    s2 = await run_scenario_2_ezptc_no_chaining()
    s3 = await run_scenario_3_ezptc_with_chaining()

    generate_markdown_report(s1, s2, s3)


def _format_code_for_report(code_blocks, scenario_type):
    """Format code blocks for the markdown report."""
    if not code_blocks:
        return "_No code generated._"
    parts = []
    for i, block in enumerate(code_blocks, 1):
        if len(code_blocks) > 1:
            parts.append(f"**Turn {i}:**")
        parts.append(f"```python\n{block}\n```")
    return "\n\n".join(parts)


def generate_markdown_report(s1, s2, s3):
    """Generate a detailed markdown comparison report."""

    s1_code = _format_code_for_report(s1["code_blocks"], "traditional")
    s2_code = _format_code_for_report(s2["code_blocks"], "ezptc")
    s3_code = _format_code_for_report(s3["code_blocks"], "ezptc")

    report = """
# Three-Way Comparison: Tool Calling Approaches

**Query:** "{query}"

**Date:** {date}

## Summary Table

| Metric | Traditional | ez-ptc (no chaining) | ez-ptc (with chaining) | Best |
|--------|-------------|----------------------|------------------------|------|
| **LLM Turns** | {s1_turns} | {s2_turns} | {s3_turns} | {best_turns} |
| **Tool Calls** | {s1_tools} | {s2_tools} | {s3_tools} | {best_tools} |
| **Total Tokens** | {s1_tokens:,} | {s2_tokens:,} | {s3_tokens:,} | {best_tokens} |
| **Input Tokens** | {s1_input:,} | {s2_input:,} | {s3_input:,} | - |
| **Output Tokens** | {s1_output:,} | {s2_output:,} | {s3_output:,} | - |
| **Execution Time** | {s1_time:.2f}s | {s2_time:.2f}s | {s3_time:.2f}s | {best_time} |

## Detailed Analysis

### Scenario 1: Traditional Tool Calling
**Approach:** Pydantic AI with individual tool functions

**How it works:**
- LLM calls each tool individually
- Each tool call requires a separate round-trip
- LLM must synthesize all results at the end

**Results:**
- LLM Turns: {s1_turns}
- Tool Calls: {s1_tools}
- Total Tokens: {s1_tokens:,}
- Time: {s1_time:.2f}s

**LLM-Generated Code:**

{s1_code}

### Scenario 2: ez-ptc WITHOUT Tool Chaining
**Approach:** ez-ptc with `assist_tool_chaining=False`

**How it works:**
- LLM writes Python code to orchestrate tools
- Toolkit doesn't document return types in detail
- Single execution of generated code

**Results:**
- LLM Turns: {s2_turns}
- Tool Calls: {s2_tools}
- Total Tokens: {s2_tokens:,}
- Time: {s2_time:.2f}s

**LLM-Generated Code:**

{s2_code}

### Scenario 3: ez-ptc WITH Tool Chaining
**Approach:** ez-ptc with `assist_tool_chaining=True`

**How it works:**
- LLM writes Python code to orchestrate tools
- Toolkit documents exact return types (enables chaining)
- LLM can filter/transform data between tool calls
- Single execution of generated code

**Results:**
- LLM Turns: {s3_turns}
- Tool Calls: {s3_tools}
- Total Tokens: {s3_tokens:,}
- Time: {s3_time:.2f}s

**LLM-Generated Code:**

{s3_code}

## Key Improvements

### ez-ptc (with chaining) vs Traditional:
- **{turns_improvement:.1f}%** fewer LLM turns
- **{tokens_improvement:.1f}%** fewer tokens
- **{time_improvement:.1f}%** faster execution

### ez-ptc (with chaining) vs ez-ptc (no chaining):
- **{chaining_turn_diff:.1f}%** {chaining_turn_dir} LLM turns
- **{chaining_token_diff:.1f}%** {chaining_token_dir} tokens
- **{chaining_time_diff:.1f}%** {chaining_time_dir} execution time

## Conclusion

**Winner:** {winner}

The key advantage of ez-ptc with tool chaining is that it:
1. Reduces LLM round-trips significantly
2. Enables programmatic filtering between tool calls
3. Returns only relevant data (selective output)
4. Saves tokens and reduces latency

Tool chaining documentation helps the LLM understand exact return types, enabling better code generation and data manipulation.
""".format(
        query=QUERY,
        date=time.strftime("%Y-%m-%d %H:%M:%S"),
        s1_code=s1_code,
        s2_code=s2_code,
        s3_code=s3_code,
        s1_turns=s1['turns'],
        s2_turns=s2['turns'],
        s3_turns=s3['turns'],
        s1_tools=s1['tool_calls'],
        s2_tools=s2['tool_calls'],
        s3_tools=s3['tool_calls'],
        s1_tokens=s1['total_tokens'],
        s2_tokens=s2['total_tokens'],
        s3_tokens=s3['total_tokens'],
        s1_input=s1['prompt_tokens'],
        s2_input=s2['prompt_tokens'],
        s3_input=s3['prompt_tokens'],
        s1_output=s1['completion_tokens'],
        s2_output=s2['completion_tokens'],
        s3_output=s3['completion_tokens'],
        s1_time=s1['time'],
        s2_time=s2['time'],
        s3_time=s3['time'],
        best_turns=f"S{1 if s1['turns'] <= min(s2['turns'], s3['turns']) else 2 if s2['turns'] < s3['turns'] else 3}",
        best_tools=f"S{1 if s1['tool_calls'] <= min(s2['tool_calls'], s3['tool_calls']) else 2 if s2['tool_calls'] < s3['tool_calls'] else 3}",
        best_tokens=f"S{1 if s1['total_tokens'] <= min(s2['total_tokens'], s3['total_tokens']) else 2 if s2['total_tokens'] < s3['total_tokens'] else 3}",
        best_time=f"S{1 if s1['time'] <= min(s2['time'], s3['time']) else 2 if s2['time'] < s3['time'] else 3}",
        turns_improvement=((s1['turns'] - s3['turns']) / s1['turns'] * 100) if s1['turns'] > 0 else 0,
        tokens_improvement=((s1['total_tokens'] - s3['total_tokens']) / s1['total_tokens'] * 100) if s1['total_tokens'] > 0 else 0,
        time_improvement=((s1['time'] - s3['time']) / s1['time'] * 100) if s1['time'] > 0 else 0,
        chaining_turn_diff=abs((s2['turns'] - s3['turns']) / s2['turns'] * 100) if s2['turns'] > 0 else 0,
        chaining_turn_dir="fewer" if s3['turns'] < s2['turns'] else "more",
        chaining_token_diff=abs((s2['total_tokens'] - s3['total_tokens']) / s2['total_tokens'] * 100) if s2['total_tokens'] > 0 else 0,
        chaining_token_dir="fewer" if s3['total_tokens'] < s2['total_tokens'] else "more",
        chaining_time_diff=abs((s2['time'] - s3['time']) / s2['time'] * 100) if s2['time'] > 0 else 0,
        chaining_time_dir="faster" if s3['time'] < s2['time'] else "slower",
        winner="Scenario 3 (ez-ptc with tool chaining)" if s3['total_tokens'] < min(s1['total_tokens'], s2['total_tokens']) else "Scenario 2 (ez-ptc no chaining)" if s2['total_tokens'] < s1['total_tokens'] else "Scenario 1 (Traditional)"
    )

    with open('./COMPARISON_REPORT.md', 'w') as f:
        f.write(report)

    print("\n" + "="*80)
    print("Comparison report saved to: COMPARISON_REPORT.md")
    print("="*80 + "\n")


if __name__ == "__main__":
    asyncio.run(main())
