"""tau-bench integration runner.

Executes tau-bench evaluations and captures trajectories for MangoDesk.
"""

import os
import sys
import time
from pathlib import Path
from typing import Any

from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
from rich.table import Table

from mangodesk.config import get_openai_api_key, get_anthropic_api_key
from mangodesk.core.api import MangoAPIClient

console = Console()


def setup_api_keys(model_provider: str) -> None:
    """Set up API keys in environment for tau-bench."""
    if model_provider == "openai":
        api_key = get_openai_api_key()
        if api_key:
            os.environ["OPENAI_API_KEY"] = api_key
    elif model_provider == "anthropic":
        api_key = get_anthropic_api_key()
        if api_key:
            os.environ["ANTHROPIC_API_KEY"] = api_key


def run_tau_bench_evaluation(
    tau_bench_path: Path,
    env: str,
    model: str,
    model_provider: str,
    agent_strategy: str,
    user_model: str,
    user_model_provider: str,
    user_strategy: str,
    max_concurrency: int,
    task_ids: list[int] | None,
    trials: int,
    upload_results: bool,
    api_url: str | None,
    api_key: str | None,
) -> None:
    """Run tau-bench evaluation with MangoDesk tracking.

    Args:
        tau_bench_path: Path to tau-bench installation
        env: Environment name (retail, airline, telecom)
        model: Model name for agent
        model_provider: Model provider (openai, anthropic, etc.)
        agent_strategy: Agent strategy (tool-calling, react, act)
        user_model: Model for user simulation
        user_model_provider: Provider for user model
        user_strategy: User strategy (llm, react, verify, reflection)
        max_concurrency: Maximum concurrent tasks
        task_ids: Specific task IDs to run (None for all)
        trials: Number of trials per task
        upload_results: Whether to upload to MangoDesk
        api_url: MangoDesk API URL
        api_key: MangoDesk API key
    """
    # Set up API keys in environment
    setup_api_keys(model_provider)
    setup_api_keys(user_model_provider)

    # Add tau-bench to Python path if needed
    if str(tau_bench_path) not in sys.path:
        sys.path.insert(0, str(tau_bench_path))

    console.print("[bold]MangoDesk Evaluation (tau-bench)[/bold]")
    console.print()
    console.print(f"  Environment:     {env}")
    console.print(f"  Model:           {model} ({model_provider})")
    console.print(f"  Agent Strategy:  {agent_strategy}")
    console.print(f"  User Model:      {user_model} ({user_model_provider})")
    console.print(f"  User Strategy:   {user_strategy}")
    console.print(f"  Concurrency:     {max_concurrency}")
    console.print(f"  Trials:          {trials}")
    if task_ids:
        console.print(f"  Task IDs:        {task_ids}")
    console.print()

    start_time = time.time()

    try:
        # Import tau-bench modules
        from tau_bench.run import run
        from tau_bench.types import RunConfig

        # Build the run configuration
        config = RunConfig(
            env=env,
            model=model,
            model_provider=model_provider,
            user_model=user_model,
            user_model_provider=user_model_provider,
            user_strategy=user_strategy,
            agent_strategy=agent_strategy,
            max_concurrency=max_concurrency,
            task_ids=task_ids,
            n=trials,
        )

        console.print("Running evaluation...")
        console.print()

        # Run tau-bench
        results = run(config)

        duration = time.time() - start_time

        # Process and display results
        display_results(results, model, env, duration)

        # Upload to MangoDesk if configured
        if upload_results and api_key and api_url:
            upload_to_mangodesk(
                results=results,
                env=env,
                model=model,
                model_provider=model_provider,
                agent_strategy=agent_strategy,
                user_model=user_model,
                user_model_provider=user_model_provider,
                user_strategy=user_strategy,
                trials=trials,
                api_url=api_url,
                api_key=api_key,
            )

    except ImportError as e:
        console.print(f"[red]Error importing tau-bench:[/red] {e}")
        console.print()
        console.print("Make sure tau-bench is installed:")
        console.print("  mangodesk init")
        console.print()
        console.print("Or specify the path:")
        console.print("  mangodesk eval --tau-bench-path /path/to/tau-bench ...")
        raise SystemExit(1)

    except Exception as e:
        console.print(f"[red]Error during evaluation:[/red] {e}")
        raise SystemExit(1)


def display_results(results: Any, model: str, env: str, duration: float) -> None:
    """Display evaluation results in a table format.

    Args:
        results: Results from tau-bench run
        model: Model name
        env: Environment name
        duration: Total evaluation duration in seconds
    """
    # Create results table
    table = Table(title=f"Evaluation Results - {env} / {model}")

    table.add_column("Task ID", style="cyan")
    table.add_column("Status", justify="center")
    table.add_column("Reward", justify="right")
    table.add_column("Steps", justify="right")

    # Parse results from tau-bench
    # tau-bench returns a list of TaskResult objects
    passed_count = 0
    total_count = 0
    total_reward = 0.0

    if hasattr(results, "__iter__"):
        for result in results:
            total_count += 1

            # Extract data from tau-bench result
            task_id = getattr(result, "task_id", "?")
            reward = getattr(result, "reward", 0.0)
            passed = reward >= 1.0
            steps = len(getattr(result, "traj", [])) if hasattr(result, "traj") else 0

            if passed:
                passed_count += 1
            total_reward += reward

            status = "[green]PASS[/green]" if passed else "[red]FAIL[/red]"
            table.add_row(str(task_id), status, f"{reward:.2f}", str(steps))

    console.print(table)
    console.print()

    # Summary
    pass_rate = (passed_count / total_count * 100) if total_count > 0 else 0
    avg_reward = (total_reward / total_count) if total_count > 0 else 0

    if passed_count == total_count and total_count > 0:
        console.print(f"[green bold]All {total_count} tasks passed![/green bold]")
    else:
        console.print(f"[yellow]{passed_count}/{total_count} passed ({pass_rate:.1f}%)[/yellow]")

    console.print(f"Average reward: {avg_reward:.3f}")
    console.print(f"Duration: {duration:.1f}s")


def upload_to_mangodesk(
    results: Any,
    env: str,
    model: str,
    model_provider: str,
    agent_strategy: str,
    user_model: str,
    user_model_provider: str,
    user_strategy: str,
    trials: int,
    api_url: str,
    api_key: str,
) -> None:
    """Upload evaluation results to MangoDesk cloud.

    Args:
        results: tau-bench results
        env: Environment name
        model: Model name
        model_provider: Model provider
        agent_strategy: Agent strategy used
        user_model: User model
        user_model_provider: User model provider
        user_strategy: User strategy used
        trials: Number of trials
        api_url: MangoDesk API URL
        api_key: MangoDesk API key
    """
    console.print()
    console.print("Uploading results to MangoDesk...")

    client = MangoAPIClient(api_url, api_key)

    # Convert tau-bench results to MangoDesk rollout format
    rollouts = []

    if hasattr(results, "__iter__"):
        for result in results:
            task_id = str(getattr(result, "task_id", "unknown"))
            reward = getattr(result, "reward", 0.0)
            passed = reward >= 1.0
            traj = getattr(result, "traj", [])

            # Convert trajectory to trace format
            trace = convert_trajectory_to_trace(traj)

            # Get prompt from task info if available
            # tau-bench stores the instruction in different places depending on version
            prompt = ""
            # Try multiple possible paths for the task instruction/prompt
            if hasattr(result, "task"):
                task = result.task
                if hasattr(task, "instruction"):
                    prompt = task.instruction
                elif hasattr(task, "prompt"):
                    prompt = task.prompt
                elif hasattr(task, "user_instruction"):
                    prompt = task.user_instruction
            # Fall back to result-level attributes
            if not prompt and hasattr(result, "instruction"):
                prompt = result.instruction
            if not prompt and hasattr(result, "prompt"):
                prompt = result.prompt
            # Try to extract from trajectory if still empty
            if not prompt and traj:
                for item in traj:
                    if hasattr(item, "role") and item.role == "user":
                        content = getattr(item, "content", "")
                        if content:
                            prompt = content
                            break

            rollouts.append({
                "task_id": task_id,
                "prompt": prompt,
                "trace": trace,
                "passed": passed,
                "score": reward,
                "duration_ms": 0,  # tau-bench doesn't track this per-task
                "token_count": 0,  # tau-bench doesn't track this
                "failure_reason": None if passed else "Task not completed successfully",
                # tau-bench specific fields
                "agent_strategy": agent_strategy,
                "user_strategy": user_strategy,
                "user_model": f"{user_model} ({user_model_provider})",
            })

    try:
        # Add tau-bench metadata to model name
        model_with_strategy = f"{model} ({agent_strategy})"

        response = client.create_run(
            environment_name=env,
            model=model_with_strategy,
            rollouts=rollouts,
        )

        dashboard_url = f"{api_url}{response['dashboard_url']}"
        console.print("[green]✓[/green] Uploaded!")
        console.print()
        console.print(f"[bold]View results:[/bold] {dashboard_url}")

    except Exception as e:
        console.print(f"[red]Failed to upload:[/red] {e}")


def convert_trajectory_to_trace(traj: list) -> list[dict]:
    """Convert tau-bench trajectory to MangoDesk trace format.

    Args:
        traj: tau-bench trajectory (list of actions/observations)

    Returns:
        MangoDesk trace format (list of messages with roles)
    """
    trace = []

    for item in traj:
        if hasattr(item, "role"):
            # Already has role attribute
            role = item.role
            content = getattr(item, "content", str(item))

            entry = {"role": role, "content": content}

            # Handle tool calls
            if hasattr(item, "tool_calls") and item.tool_calls:
                entry["tool_calls"] = [
                    {
                        "name": tc.function.name if hasattr(tc, "function") else str(tc),
                        "arguments": tc.function.arguments if hasattr(tc, "function") else {},
                    }
                    for tc in item.tool_calls
                ]

            trace.append(entry)
        else:
            # Generic item - convert to string
            trace.append({
                "role": "system",
                "content": str(item),
            })

    return trace


def get_environment_task_ids(tau_bench_path: Path, env: str) -> list[int]:
    """Get the list of task IDs available in a tau-bench environment.
    
    Args:
        tau_bench_path: Path to tau-bench installation
        env: Environment name (retail, airline, etc.)
        
    Returns:
        List of available task IDs
    """
    # Add tau-bench to path temporarily
    if str(tau_bench_path) not in sys.path:
        sys.path.insert(0, str(tau_bench_path))
    
    try:
        # Import the environment module to get task count
        if env == "retail":
            from tau_bench.envs.retail import TASKS
        elif env == "airline":
            from tau_bench.envs.airline import TASKS
        else:
            # Try generic import pattern
            try:
                import importlib
                env_module = importlib.import_module(f"tau_bench.envs.{env}")
                TASKS = getattr(env_module, "TASKS", None)
                if TASKS is None:
                    # Fall back to checking for tasks attribute
                    TASKS = getattr(env_module, "tasks", [])
            except ImportError:
                console.print(f"[yellow]Warning:[/yellow] Could not determine task count for environment '{env}'")
                console.print("Please specify task IDs explicitly with --task-ids")
                raise SystemExit(1)
        
        # TASKS is typically a list, return indices
        return list(range(len(TASKS)))
        
    except ImportError as e:
        console.print(f"[yellow]Warning:[/yellow] Could not import environment '{env}': {e}")
        console.print("Please specify task IDs explicitly with --task-ids")
        raise SystemExit(1)


def run_remote_evaluation(
    tau_bench_path: Path,
    env: str,
    model: str,
    model_provider: str,
    agent_strategy: str,
    user_model: str,
    user_model_provider: str,
    user_strategy: str,
    task_ids: list[int] | None,
    trials: int,
    upload_results: bool,
    api_url: str,
    api_key: str,
) -> None:
    """Run tau-bench evaluation on remote workers.

    Packages the local tau-bench directory, uploads it to the API,
    and polls for results as remote workers process the tasks.

    Args:
        tau_bench_path: Path to tau-bench installation (will be packaged)
        env: Environment name (retail, airline, telecom)
        model: Model name for agent
        model_provider: Model provider (openai, anthropic, etc.)
        agent_strategy: Agent strategy (tool-calling, react, act)
        user_model: Model for user simulation
        user_model_provider: Provider for user model
        user_strategy: User strategy (llm, react, verify, reflection)
        task_ids: Specific task IDs to run (None for all tasks in env)
        trials: Number of trials per task
        upload_results: Whether to upload to MangoDesk after completion
        api_url: MangoDesk API URL
        api_key: MangoDesk API key
    """
    import json
    import httpx
    from mangodesk.core.bundle import create_bundle, get_bundle_size_mb

    # Determine task IDs if not specified
    if task_ids is None:
        console.print(f"Determining tasks in environment '{env}'...")
        task_ids = get_environment_task_ids(tau_bench_path, env)
        console.print(f"  Found {len(task_ids)} tasks")
        console.print()

    console.print("[bold]MangoDesk Remote Evaluation (tau-bench)[/bold]")
    console.print()
    console.print(f"  Environment:     {env}")
    console.print(f"  Model:           {model} ({model_provider})")
    console.print(f"  Agent Strategy:  {agent_strategy}")
    console.print(f"  User Model:      {user_model} ({user_model_provider})")
    console.print(f"  User Strategy:   {user_strategy}")
    console.print(f"  Trials:          {trials}")
    console.print(f"  Mode:            [cyan]REMOTE[/cyan] (parallel execution)")
    console.print(f"  Tasks:           {len(task_ids)} tasks")
    console.print()

    start_time = time.time()

    # Package tau-bench directory
    console.print("Packaging tau-bench code...")
    try:
        bundle_data = create_bundle(tau_bench_path)
        bundle_size = get_bundle_size_mb(bundle_data)
        console.print(f"  Bundle size: {bundle_size:.1f} MB")
    except Exception as e:
        console.print(f"[red]Error packaging tau-bench:[/red] {e}")
        raise SystemExit(1)

    # Collect API keys from local environment
    api_keys = {}
    if os.environ.get("OPENAI_API_KEY"):
        api_keys["OPENAI_API_KEY"] = os.environ["OPENAI_API_KEY"]
    if os.environ.get("ANTHROPIC_API_KEY"):
        api_keys["ANTHROPIC_API_KEY"] = os.environ["ANTHROPIC_API_KEY"]
    if os.environ.get("GOOGLE_API_KEY"):
        api_keys["GOOGLE_API_KEY"] = os.environ["GOOGLE_API_KEY"]
    if os.environ.get("MISTRAL_API_KEY"):
        api_keys["MISTRAL_API_KEY"] = os.environ["MISTRAL_API_KEY"]
    if os.environ.get("COHERE_API_KEY"):
        api_keys["COHERE_API_KEY"] = os.environ["COHERE_API_KEY"]
    if os.environ.get("TOGETHER_API_KEY"):
        api_keys["TOGETHER_API_KEY"] = os.environ["TOGETHER_API_KEY"]
    
    # Validate required API key is present for the model provider
    provider_key_map = {
        "openai": "OPENAI_API_KEY",
        "anthropic": "ANTHROPIC_API_KEY",
        "google": "GOOGLE_API_KEY",
        "mistral": "MISTRAL_API_KEY",
        "cohere": "COHERE_API_KEY",
        "together": "TOGETHER_API_KEY",
    }
    required_key = provider_key_map.get(model_provider)
    if required_key and required_key not in api_keys:
        console.print(f"[red]Error:[/red] {required_key} environment variable is required for {model_provider} provider")
        console.print(f"Please set it before running: export {required_key}=your-key")
        raise SystemExit(1)
    
    # Prepare config
    config = {
        "env": env,
        "model": model,
        "model_provider": model_provider,
        "task_ids": task_ids,
        "config": {
            "agent_strategy": agent_strategy,
            "user_model": user_model,
            "user_model_provider": user_model_provider,
            "user_strategy": user_strategy,
            "trials": trials,
        },
        "api_keys": api_keys if api_keys else None,
    }

    # Submit to API
    console.print()
    console.print("Submitting evaluation to remote workers...")
    
    try:
        client = httpx.Client(
            timeout=300.0,
            headers={"Authorization": f"Bearer {api_key}"},
        )
        
        response = client.post(
            f"{api_url}/v1/remote-eval",
            data={"config": json.dumps(config)},
            files={"code_bundle": ("bundle.tar.gz", bundle_data, "application/gzip")},
        )
        response.raise_for_status()
        eval_data = response.json()
        
        eval_id = eval_data["id"]
        total_tasks = eval_data["total_tasks"]
        
        console.print(f"[green]✓[/green] Submitted! Eval ID: {eval_id}")
        console.print(f"  Total tasks: {total_tasks}")
        console.print()
        
    except httpx.HTTPError as e:
        console.print(f"[red]Error submitting evaluation:[/red] {e}")
        if hasattr(e, "response") and e.response is not None:
            console.print(e.response.text)
        raise SystemExit(1)

    # Poll for results
    console.print("Waiting for remote workers...")
    console.print()
    
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        console=console,
    ) as progress:
        task = progress.add_task("Processing tasks...", total=total_tasks)
        
        last_completed = 0
        poll_interval = 2  # seconds
        
        while True:
            try:
                status_response = client.get(
                    f"{api_url}/v1/remote-eval/{eval_id}/status",
                )
                status_response.raise_for_status()
                status_data = status_response.json()
                
                completed = status_data["completed_tasks"]
                passed = status_data["passed_tasks"]
                status = status_data["status"]
                
                # Update progress
                if completed > last_completed:
                    progress.update(task, completed=completed)
                    last_completed = completed
                
                # Update description
                progress.update(
                    task, 
                    description=f"Processing tasks... ({passed}/{completed} passed)"
                )
                
                # Check if done
                if status in ("completed", "failed"):
                    progress.update(task, completed=total_tasks)
                    break
                
                time.sleep(poll_interval)
                
            except httpx.HTTPError as e:
                console.print(f"[yellow]Warning: Poll error:[/yellow] {e}")
                time.sleep(poll_interval)

    duration = time.time() - start_time

    # Fetch final results
    try:
        final_response = client.get(f"{api_url}/v1/remote-eval/{eval_id}/status")
        final_response.raise_for_status()
        final_data = final_response.json()
    except httpx.HTTPError as e:
        console.print(f"[red]Error fetching final results:[/red] {e}")
        raise SystemExit(1)

    # Display results (same format as local)
    display_remote_results(final_data, model, env, duration)

    # Upload to MangoDesk if configured
    if upload_results:
        upload_remote_to_mangodesk(
            results=final_data,
            env=env,
            model=model,
            model_provider=model_provider,
            agent_strategy=agent_strategy,
            user_model=user_model,
            user_model_provider=user_model_provider,
            user_strategy=user_strategy,
            trials=trials,
            api_url=api_url,
            api_key=api_key,
        )


def display_remote_results(data: dict, model: str, env: str, duration: float) -> None:
    """Display remote evaluation results in a table format.

    Args:
        data: Status data from remote eval API
        model: Model name
        env: Environment name
        duration: Total evaluation duration in seconds
    """
    table = Table(title=f"Evaluation Results - {env} / {model}")

    table.add_column("Task ID", style="cyan")
    table.add_column("Status", justify="center")
    table.add_column("Reward", justify="right")
    table.add_column("Steps", justify="right")

    passed_count = 0
    total_count = 0
    total_reward = 0.0

    for task in data.get("tasks", []):
        total_count += 1
        
        task_id = task["task_id"]
        reward = task.get("reward", 0.0) or 0.0
        passed = task.get("passed", False)
        trace = task.get("trace", [])
        steps = len(trace) if trace else 0
        
        if passed:
            passed_count += 1
        total_reward += reward
        
        if task["status"] == "failed":
            status = "[red]ERROR[/red]"
        elif passed:
            status = "[green]PASS[/green]"
        else:
            status = "[red]FAIL[/red]"
        
        table.add_row(str(task_id), status, f"{reward:.2f}", str(steps))

    console.print(table)
    console.print()

    # Summary
    pass_rate = (passed_count / total_count * 100) if total_count > 0 else 0
    avg_reward = (total_reward / total_count) if total_count > 0 else 0

    if passed_count == total_count and total_count > 0:
        console.print(f"[green bold]All {total_count} tasks passed![/green bold]")
    else:
        console.print(f"[yellow]{passed_count}/{total_count} passed ({pass_rate:.1f}%)[/yellow]")

    console.print(f"Average reward: {avg_reward:.3f}")
    console.print(f"Duration: {duration:.1f}s")


def upload_remote_to_mangodesk(
    results: dict,
    env: str,
    model: str,
    model_provider: str,
    agent_strategy: str,
    user_model: str,
    user_model_provider: str,
    user_strategy: str,
    trials: int,
    api_url: str,
    api_key: str,
) -> None:
    """Upload remote evaluation results to MangoDesk cloud.

    Args:
        results: Remote eval status data
        env: Environment name
        model: Model name
        model_provider: Model provider
        agent_strategy: Agent strategy used
        user_model: User model
        user_model_provider: User model provider
        user_strategy: User strategy used
        trials: Number of trials
        api_url: MangoDesk API URL
        api_key: MangoDesk API key
    """
    console.print()
    console.print("Uploading results to MangoDesk...")

    client = MangoAPIClient(api_url, api_key)

    # Convert remote results to MangoDesk rollout format
    rollouts = []

    for task in results.get("tasks", []):
        task_id = str(task["task_id"])
        reward = task.get("reward", 0.0) or 0.0
        passed = task.get("passed", False)
        trace = task.get("trace", []) or []
        prompt = task.get("prompt", "")
        error = task.get("error")

        rollouts.append({
            "task_id": task_id,
            "prompt": prompt,
            "trace": trace,
            "passed": passed,
            "score": reward,
            "duration_ms": 0,
            "token_count": 0,
            "failure_reason": error if not passed else None,
            "agent_strategy": agent_strategy,
            "user_strategy": user_strategy,
            "user_model": f"{user_model} ({user_model_provider})",
        })

    try:
        model_with_strategy = f"{model} ({agent_strategy})"

        response = client.create_run(
            environment_name=env,
            model=model_with_strategy,
            rollouts=rollouts,
        )

        dashboard_url = f"{api_url}{response['dashboard_url']}"
        console.print("[green]✓[/green] Uploaded!")
        console.print()
        console.print(f"[bold]View results:[/bold] {dashboard_url}")

    except Exception as e:
        console.print(f"[red]Failed to upload:[/red] {e}")
