import argparse
import asyncio
import json
import sys
import os
from datetime import datetime
from typing import Any

import multillm


# Built-in tools for chat providers
BUILTIN_TOOLS = {
    "get_current_time": {
        "type": "function",
        "function": {
            "name": "get_current_time",
            "description": "Get the current date and time",
            "parameters": {
                "type": "object",
                "properties": {
                    "timezone": {
                        "type": "string",
                        "description": "Timezone (e.g., UTC, EST). Optional, defaults to local time."
                    }
                }
            }
        }
    },
    "calculate": {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "Perform a mathematical calculation",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "Mathematical expression to evaluate (e.g., '2+2', '15*23')"
                    }
                },
                "required": ["expression"]
            }
        }
    },
    "get_weather": {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get weather information for a location (mock data)",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "City name or location"
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                        "description": "Temperature unit"
                    }
                },
                "required": ["location"]
            }
        }
    },
    "ask_user": {
        "type": "function",
        "function": {
            "name": "ask_user",
            "description": "Ask the user a question and get their response. Use this when you need user input or clarification.",
            "parameters": {
                "type": "object",
                "properties": {
                    "question": {
                        "type": "string",
                        "description": "The question to ask the user"
                    },
                    "options": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Optional list of suggested answers (user can still provide their own)"
                    }
                },
                "required": ["question"]
            }
        }
    }
}


# Tool implementations
def get_current_time(timezone: str = None) -> dict:
    """Get current time."""
    now = datetime.now()
    return {
        "time": now.strftime("%Y-%m-%d %H:%M:%S"),
        "timezone": timezone or "local",
        "timestamp": now.timestamp()
    }


def calculate(expression: str) -> dict:
    """Safely evaluate a mathematical expression."""
    try:
        # Only allow basic math operations
        allowed_chars = set("0123456789+-*/()%. ")
        if not all(c in allowed_chars for c in expression):
            return {"error": "Invalid characters in expression"}

        result = eval(expression, {"__builtins__": {}}, {})
        return {
            "expression": expression,
            "result": result
        }
    except Exception as e:
        return {
            "error": str(e),
            "expression": expression
        }


def get_weather(location: str, unit: str = "celsius") -> dict:
    """Mock weather data."""
    # Mock data - in real implementation would call weather API
    import random
    temp_c = random.randint(15, 30)
    temp_f = int(temp_c * 9/5 + 32)

    conditions = ["sunny", "cloudy", "rainy", "partly cloudy"]

    return {
        "location": location,
        "temperature": temp_c if unit == "celsius" else temp_f,
        "unit": unit,
        "condition": random.choice(conditions),
        "humidity": random.randint(30, 80)
    }


def ask_user(question: str, options: list[str] = None) -> dict:
    """
    Ask the user a question and collect their response.

    This is an interactive tool that displays a question to the user
    and waits for their input.
    """
    print("\n" + "=" * 70, file=sys.stderr)
    print("❓ QUESTION FROM ASSISTANT", file=sys.stderr)
    print("=" * 70, file=sys.stderr)
    print(f"\n{question}\n", file=sys.stderr)

    if options:
        print("Suggested options:", file=sys.stderr)
        for i, opt in enumerate(options, 1):
            print(f"  {i}. {opt}", file=sys.stderr)
        print("\nYou can select a number or provide your own answer.", file=sys.stderr)

    print("\nYour answer: ", file=sys.stderr, end="", flush=True)

    try:
        # Read from stdin
        answer = input()

        # If user entered a number and we have options, use that option
        if options and answer.strip().isdigit():
            idx = int(answer.strip()) - 1
            if 0 <= idx < len(options):
                answer = options[idx]

        print("=" * 70 + "\n", file=sys.stderr)

        return {
            "question": question,
            "answer": answer,
            "selected_from_options": answer in options if options else False
        }

    except (EOFError, KeyboardInterrupt):
        print("\n", file=sys.stderr)
        print("=" * 70 + "\n", file=sys.stderr)
        return {
            "question": question,
            "answer": None,
            "error": "User cancelled input"
        }


TOOL_FUNCTIONS = {
    "get_current_time": get_current_time,
    "calculate": calculate,
    "get_weather": get_weather,
    "ask_user": ask_user,
}


async def run_with_tools(
    client: multillm.Client,
    model: str,
    prompt: str,
    tools: list[dict],
    verbose: bool = False
) -> str:
    """
    Run a chat completion with tool support.

    Implements the full tool calling loop:
    1. Send request with tools
    2. Check for tool calls
    3. Execute tools
    4. Send results back
    5. Get final response
    """
    messages = [{"role": "user", "content": prompt}]

    max_iterations = 5  # Prevent infinite loops
    iteration = 0

    while iteration < max_iterations:
        iteration += 1

        response = await client.chat_complete(model, messages, tools=tools)

        # Check for tool calls
        if not response.choices[0].message.tool_calls:
            # No tool calls, we have the final response
            return response.text

        # Process tool calls
        tool_calls = response.choices[0].message.tool_calls

        # Show that we're using tools
        print(f"\n[Using {len(tool_calls)} tool(s)]")

        # Add assistant message with tool calls
        messages.append({
            "role": "assistant",
            "content": response.text or "",
            "tool_calls": tool_calls
        })

        # Execute each tool call
        for tool_call in tool_calls:
            function_name = tool_call["function"]["name"]
            function_args = tool_call["function"]["arguments"]

            # Parse arguments if string
            if isinstance(function_args, str):
                function_args = json.loads(function_args)

            # Show tool call (always show function name)
            if verbose:
                print(f"  → {function_name}({json.dumps(function_args)})")
            else:
                print(f"  → {function_name}()")

            # Execute the tool
            if function_name in TOOL_FUNCTIONS:
                try:
                    result = TOOL_FUNCTIONS[function_name](**function_args)
                    if verbose:
                        print(f"  ← {json.dumps(result, indent=2)}")
                    else:
                        # Show brief result summary
                        if "error" in result:
                            print(f"  ✗ Error: {result['error']}")
                        elif "result" in result:
                            print(f"  ✓ Result: {result['result']}")
                        else:
                            print(f"  ✓ Success")
                except Exception as e:
                    result = {"error": str(e)}
                    print(f"  ✗ Error: {e}")
            else:
                result = {"error": f"Unknown function: {function_name}"}
                print(f"  ✗ Unknown function")

            # Add tool result to messages
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call["id"],
                "name": function_name,
                "content": json.dumps(result)
            })

        print()  # Blank line before continuing

    # If we hit max iterations, return what we have
    return "Maximum tool calling iterations reached"


async def run_agentic(
    model: str,
    prompt: str,
    tools: list[multillm.Tool] | None = None,
    options: multillm.AgentOptions | None = None,
    verbose: bool = False
) -> str:
    """
    Run a query using the agentic API.

    Uses agentwrap for chat providers, native agent API for agent providers.
    """
    client = multillm.Client()

    # For Claude, if AskUserQuestion is requested, provide custom ask_user tool instead
    provider_name = model.split("/")[0]
    if provider_name == "claude" and options and options.allowed_tools:
        if "AskUserQuestion" in options.allowed_tools:
            # Remove AskUserQuestion (SDK built-in doesn't work interactively)
            options.allowed_tools = [t for t in options.allowed_tools if t != "AskUserQuestion"]

            # Add our custom ask_user tool
            if not tools:
                tools = []

            # Create ask_user tool for Claude
            ask_user_claude = multillm.Tool(
                name="ask_user",
                description="Ask the user a question and get their response. Use this when you need user input or clarification.",
                parameters={
                    "type": "object",
                    "properties": {
                        "question": {
                            "type": "string",
                            "description": "The question to ask the user"
                        },
                        "options": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "Optional suggested answers"
                        }
                    },
                    "required": ["question"]
                },
                handler=ask_user  # Use the same handler as chat providers
            )
            tools.append(ask_user_claude)

            print("ℹ️  Using custom 'ask_user' tool instead of AskUserQuestion for interactive prompting", file=sys.stderr)

    # Collect text responses
    text_parts = []
    tool_uses = []

    async for msg in client.run(model, prompt, options=options, tools=tools):
        if msg.type == "text":
            text_parts.append(msg.content)
        elif msg.type == "tool_use":
            tool_uses.append(msg)
            if verbose:
                print(f"  → {msg.tool_name}({json.dumps(msg.tool_input)})", file=sys.stderr)
            else:
                print(f"  → {msg.tool_name}", file=sys.stderr)
        elif msg.type == "tool_result":
            if verbose:
                print(f"  ← {msg.tool_result}", file=sys.stderr)

    # Show tool usage summary if any tools were used
    if tool_uses and not verbose:
        print(f"\n[Used {len(tool_uses)} tool(s)]\n", file=sys.stderr)

    return " ".join(text_parts)


async def run_with_chat_tools(
    model: str,
    prompt: str,
    enabled_tools: list[str],
    verbose: bool = False
) -> str:
    """
    Run with chat provider tools using agentwrap.

    Converts built-in tools to Tool objects and uses agentwrap for execution.
    """
    # Build Tool objects from enabled tools
    tool_objects = []
    for name in enabled_tools:
        if name in BUILTIN_TOOLS:
            tool_def = BUILTIN_TOOLS[name]
            tool_objects.append(multillm.Tool(
                name=tool_def["function"]["name"],
                description=tool_def["function"]["description"],
                parameters=tool_def["function"]["parameters"],
                handler=TOOL_FUNCTIONS[name]
            ))

    if not tool_objects:
        # No valid tools, run without tools
        return await run_agentic(f"agentwrap/{model}", prompt, verbose=verbose)

    # Run with agentwrap and tools
    return await run_agentic(
        f"agentwrap/{model}",
        prompt,
        tools=tool_objects,
        verbose=verbose
    )


def main():
    """CLI entry point."""
    parser = argparse.ArgumentParser(
        description="Run prompts against LLM providers",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Chat providers (simple queries) - uses agentwrap internally
  multillm -m openai/gpt-4o -p "What is 2+2?"
  multillm -m anthropic/claude-sonnet-4-20250514 -p "Explain async/await"
  multillm -m gemini/gemini-2.0-flash-exp -p "What is Python?"

  # With built-in tools (for chat providers) - uses agentwrap with tool execution
  multillm -m openai/gpt-4o -p "What time is it?" --use-tools get_current_time
  multillm -m openai/gpt-4o -p "Calculate 15 * 23" --use-tools calculate
  multillm -m openai/gpt-4o -p "What's the weather in Tokyo?" --use-tools get_weather
  multillm -m openai/gpt-4o -p "What's 5+5 and the current time?" --use-tools calculate get_current_time

  # Interactive tools (ask user questions)
  multillm -m openai/gpt-4o -p "Ask me about my preferences and create a summary" --use-tools ask_user

  # Native agent providers (Claude with built-in tools)
  multillm -m claude/default -p "What Python version?" --allowed-tools Bash
  multillm -m claude/default -p "List files" --allowed-tools Bash Glob --max-turns 5
  multillm -m claude/default -p "Read README.md" --allowed-tools Read

  # With stdin
  cat file.txt | multillm -m openai/gpt-4o -p "Summarize:" --with-stdin

  # Permission modes (for native agents)
  multillm -m claude/default -p "Create hello.py" --allowed-tools Write --permission-mode acceptEdits

  # Verbose mode
  multillm -m openai/gpt-4o -p "Calculate 5*5" --use-tools calculate --verbose

Note:
  - Chat providers (openai, google, anthropic, etc.) are automatically wrapped with
    agentic capabilities using the 'agentwrap' provider
  - Native agent providers (claude) use their built-in agentic features
  - Use --use-tools for chat providers, --allowed-tools for native agents

Available Built-in Tools (for chat providers with --use-tools):
  get_current_time    Get current date and time
  calculate           Perform mathematical calculations
  get_weather         Get weather information (mock data)
  ask_user            Ask the user a question and get their response (interactive)

Available Tools (for agent providers with --allowed-tools):
  Read, Write, Edit, Bash, Glob, Grep, Task, WebFetch, WebSearch,
  NotebookEdit, AskUserQuestion, KillShell, EnterPlanMode, ExitPlanMode

  Common tools: Read Write Bash Glob Grep
        """,
    )

    parser.add_argument(
        "-m",
        "--model",
        required=True,
        help="Model to use (format: provider/model-name)",
    )
    parser.add_argument(
        "-p",
        "--prompt",
        required=True,
        help="Prompt to send to the model",
    )
    parser.add_argument(
        "--with-stdin",
        action="store_true",
        help="Append stdin to the prompt after a separator",
    )
    parser.add_argument(
        "--use-tools",
        nargs="+",
        choices=list(BUILTIN_TOOLS.keys()),
        help="Enable built-in tools for chat providers",
    )
    parser.add_argument(
        "--max-turns",
        type=int,
        help="Maximum turns for agent providers",
    )
    parser.add_argument(
        "--allowed-tools",
        nargs="+",
        help="Allowed tools for agent providers (e.g., Read Write Bash Glob Grep Edit)",
    )
    parser.add_argument(
        "--permission-mode",
        choices=["acceptEdits", "acceptAll", "prompt"],
        help="Permission mode for agent providers",
    )
    parser.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="Show tool execution details",
    )

    args = parser.parse_args()

    # Build the prompt, appending stdin if requested
    prompt = args.prompt
    if args.with_stdin:
        stdin_content = sys.stdin.read()
        prompt = f"{prompt}\n--- USER STDIN BEGIN ---\n{stdin_content}"

    try:
        # Determine if this is a chat or agent provider
        provider_name = args.model.split("/")[0]
        is_agent_provider = provider_name in ["claude"]  # Native agent providers

        if args.use_tools:
            # Use tool calling workflow for chat providers with agentwrap
            result_text = asyncio.run(
                run_with_chat_tools(args.model, prompt, args.use_tools, args.verbose)
            )
            print(result_text)
        else:
            # Build agent options
            options = None
            if args.max_turns is not None or args.allowed_tools or args.permission_mode:
                options = multillm.AgentOptions(
                    max_turns=args.max_turns,
                    allowed_tools=args.allowed_tools,
                    permission_mode=args.permission_mode,
                )

            # Determine which model string to use
            if is_agent_provider:
                # Use agent provider directly (claude)
                model_to_use = args.model
            else:
                # Use agentwrap for chat providers
                model_to_use = f"agentwrap/{args.model}"

            # Run with agentic API
            result_text = asyncio.run(
                run_agentic(model_to_use, prompt, options=options, verbose=args.verbose)
            )
            print(result_text)

    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        if args.verbose:
            import traceback
            traceback.print_exc(file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()