"""Main CLI entrypoint for MangoDesk."""

import json
from pathlib import Path

import click
from rich.console import Console
from rich.table import Table

from mangodesk.config import load_config, set_config, get_api_key, get_api_url
from mangodesk.cli.eval import run_evaluation

console = Console()


@click.group()
@click.version_option()
def cli():
    """MangoDesk - Tool-use RL environment evaluation CLI."""
    pass


@cli.command()
@click.argument("key_value")
def set(key_value: str):
    """Set a configuration value.

    Usage: mangodesk set KEY=value

    Examples:
        mangodesk set MANGO_API_KEY=sk-mango-xxx
        mangodesk set API_URL=http://localhost:8000
    """
    if "=" not in key_value:
        console.print("[red]Error:[/red] Use format KEY=value")
        raise SystemExit(1)

    key, value = key_value.split("=", 1)
    key = key.strip().upper()
    value = value.strip()

    # Handle empty values as None
    if value == "":
        value = None

    set_config(key, value)
    console.print(f"[green]Set[/green] {key} = {value or '(empty)'}")


@cli.command()
def config():
    """Show current configuration."""
    cfg = load_config()

    table = Table(title="MangoDesk Configuration")
    table.add_column("Key", style="cyan")
    table.add_column("Value", style="green")

    for key, value in sorted(cfg.items()):
        display_value = value
        # Mask sensitive values
        if key.endswith("_KEY") and value:
            display_value = value[:10] + "..." + value[-4:] if len(value) > 14 else "***"
        table.add_row(key, str(display_value) if display_value else "(not set)")

    console.print(table)


@cli.command()
@click.argument("name", default="my-env")
def init(name: str):
    """Initialize a new environment scaffold.

    Creates a new directory with template files for a tau-bench compatible environment.

    Usage: mangodesk init [NAME]

    Example:
        mangodesk init airline-env
    """
    env_path = Path(name)

    if env_path.exists():
        console.print(f"[red]Error:[/red] Directory '{name}' already exists")
        raise SystemExit(1)

    env_path.mkdir(parents=True)

    # Create tools.py
    tools_content = '''"""Tool definitions for the environment.

Define tools that the agent can use to interact with the environment.
Each tool should be a function decorated with @tool.
"""

from typing import Any


def get_tools() -> list[dict[str, Any]]:
    """Return list of tool schemas in OpenAI function format."""
    return [
        {
            "type": "function",
            "function": {
                "name": "search_database",
                "description": "Search the database for records matching the query",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query"
                        },
                        "limit": {
                            "type": "integer",
                            "description": "Maximum number of results to return",
                            "default": 10
                        }
                    },
                    "required": ["query"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "update_record",
                "description": "Update a record in the database",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "record_id": {
                            "type": "string",
                            "description": "The ID of the record to update"
                        },
                        "fields": {
                            "type": "object",
                            "description": "Fields to update"
                        }
                    },
                    "required": ["record_id", "fields"]
                }
            }
        }
    ]


def execute_tool(name: str, arguments: dict[str, Any], state: dict) -> dict[str, Any]:
    """Execute a tool and return the result.

    Args:
        name: Tool name
        arguments: Tool arguments
        state: Current environment state

    Returns:
        Tool execution result
    """
    if name == "search_database":
        # Mock implementation
        return {
            "results": [
                {"id": "1", "name": "Sample Record", "query_match": arguments["query"]}
            ],
            "total": 1
        }
    elif name == "update_record":
        return {"success": True, "record_id": arguments["record_id"]}
    else:
        return {"error": f"Unknown tool: {name}"}
'''
    (env_path / "tools.py").write_text(tools_content)

    # Create data.py
    data_content = '''"""Data loader for the environment.

Provides initial state and data for task execution.
"""

from typing import Any


def get_initial_state() -> dict[str, Any]:
    """Return the initial state for task execution."""
    return {
        "database": {
            "users": [
                {"id": "u1", "name": "Alice", "email": "alice@example.com"},
                {"id": "u2", "name": "Bob", "email": "bob@example.com"},
            ],
            "orders": [
                {"id": "o1", "user_id": "u1", "total": 99.99, "status": "pending"},
                {"id": "o2", "user_id": "u2", "total": 149.99, "status": "shipped"},
            ]
        },
        "config": {
            "max_results": 100,
            "timeout_seconds": 30
        }
    }


def get_user_context() -> str:
    """Return context information shown to the agent."""
    return """You are a helpful customer service agent with access to the company database.
You can search for records and update them as needed to help customers."""
'''
    (env_path / "data.py").write_text(data_content)

    # Create tasks.json
    tasks_content = {
        "environment": name,
        "description": f"Tasks for {name} environment",
        "tasks": [
            {
                "id": "task_001",
                "prompt": "Find the user named Alice and update their email to alice.new@example.com",
                "verifier": {
                    "type": "state_check",
                    "conditions": [
                        {"path": "database.users[0].email", "equals": "alice.new@example.com"}
                    ]
                }
            },
            {
                "id": "task_002",
                "prompt": "Search for all pending orders and provide a summary",
                "verifier": {
                    "type": "output_contains",
                    "patterns": ["pending", "o1", "99.99"]
                }
            },
            {
                "id": "task_003",
                "prompt": "Update order o1 status to 'completed'",
                "verifier": {
                    "type": "state_check",
                    "conditions": [
                        {"path": "database.orders[0].status", "equals": "completed"}
                    ]
                }
            }
        ]
    }
    (env_path / "tasks.json").write_text(json.dumps(tasks_content, indent=2))

    # Create README.md
    readme_content = f"""# {name}

A tau-bench compatible environment for tool-use evaluation.

## Structure

- `tools.py` - Tool definitions and execution logic
- `data.py` - Initial state and data loader
- `tasks.json` - Task definitions with verifiers

## Usage

```bash
# Run evaluation with a model
mangodesk eval ./{name} claude-3-5-sonnet

# Run specific tasks
mangodesk eval ./{name} gpt-4 --tasks task_001,task_002
```

## Tasks

| ID | Description |
|----|-------------|
| task_001 | Find and update user email |
| task_002 | Search and summarize pending orders |
| task_003 | Update order status |

## Customization

1. Edit `tools.py` to add or modify tools
2. Edit `data.py` to change the initial state
3. Edit `tasks.json` to add or modify tasks
"""
    (env_path / "README.md").write_text(readme_content)

    console.print(f"[green]Created environment:[/green] {name}/")
    console.print("  - tools.py")
    console.print("  - data.py")
    console.print("  - tasks.json")
    console.print("  - README.md")
    console.print()
    console.print(f"[cyan]Next steps:[/cyan]")
    console.print(f"  cd {name}")
    console.print(f"  mangodesk eval . claude-3-5-sonnet")


@cli.command()
@click.argument("env_path", type=click.Path(exists=True))
@click.argument("model")
@click.option("--tasks", "-t", help="Comma-separated list of task IDs to run")
def eval(env_path: str, model: str, tasks: str | None):
    """Run evaluation on an environment.

    ENV_PATH: Path to the environment directory
    MODEL: Model name to use for evaluation (e.g., claude-3-5-sonnet, gpt-4)

    Examples:
        mangodesk eval ./my-env claude-3-5-sonnet
        mangodesk eval ./my-env gpt-4 --tasks task_001,task_002
    """
    # Check for API key
    api_key = get_api_key()
    if not api_key:
        console.print("[yellow]Warning:[/yellow] MANGO_API_KEY not set. Results won't be uploaded.")
        console.print("Run: mangodesk set MANGO_API_KEY=<your-key>")
        console.print()

    task_ids = tasks.split(",") if tasks else None

    run_evaluation(
        env_path=Path(env_path),
        model=model,
        task_ids=task_ids,
        api_url=get_api_url(),
        api_key=api_key,
    )


if __name__ == "__main__":
    cli()
