"""
Interactive chat CLI for interacting with the CRM environment via an LLM.

This example demonstrates how to build a conversational agent that uses
the CRM environment tools. The LLM can make multiple tool calls to query
and modify CRM data before providing a final answer.

Requirements:
    pip install construct-labs-crm-env openai python-dotenv

Usage:
    export CRM_AGENT_API_KEY=your-crm-api-key
    export OPENAI_API_KEY=your-openrouter-api-key
    python chat.py --help
    python chat.py
    python chat.py --question "How many companies are there?"
    python chat.py --model anthropic/claude-sonnet-4 --stream

    # Using OpenAI directly:
    python chat.py --llm-base-url https://api.openai.com/v1 --model gpt-4o
"""

from __future__ import annotations

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

try:
    from dotenv import load_dotenv
except ImportError:
    load_dotenv = None  # type: ignore[assignment, misc]

try:
    from openai import OpenAI
except ImportError:
    OpenAI = None  # type: ignore[assignment, misc]

from construct_labs_crm_env import CrmAgentEnv


def _check_dependencies() -> None:
    """Check that optional dependencies are installed."""
    missing = []
    if OpenAI is None:
        missing.append("openai")
    if load_dotenv is None:
        missing.append("python-dotenv")

    if missing:
        print(
            f"Missing dependencies: {', '.join(missing)}\n"
            f"Install with: pip install construct-labs-crm-env[chat]"
        )
        sys.exit(1)


if load_dotenv is not None:
    load_dotenv()


def print_colored(text: str, color: str) -> None:
    """Print colored text to terminal."""
    colors = {
        "red": "\033[91m",
        "green": "\033[92m",
        "yellow": "\033[93m",
        "blue": "\033[94m",
        "magenta": "\033[95m",
        "cyan": "\033[96m",
        "white": "\033[97m",
        "reset": "\033[0m",
        "bold": "\033[1m",
        "dim": "\033[2m",
    }
    print(f"{colors.get(color, '')}{text}{colors['reset']}")


def print_separator(char: str = "-", width: int = 60) -> None:
    """Print a separator line."""
    print(char * width)


def format_tool_call_compact(tool_name: str, tool_args: dict) -> str:
    """Format a tool call in compact function-call style: func_name(arg1, arg2, ...)"""
    if not tool_args:
        return f"{tool_name}()"

    arg_parts = []
    for key, value in tool_args.items():
        if isinstance(value, str):
            if len(value) > 30:
                value = value[:27] + "..."
            arg_parts.append(f'{key}="{value}"')
        elif isinstance(value, (dict, list)):
            arg_parts.append(f"{key}=<{type(value).__name__}>")
        else:
            arg_parts.append(f"{key}={value}")

    return f"{tool_name}({', '.join(arg_parts)})"


def _stream_response(
    client: OpenAI,
    model_name: str,
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]],
    print_output: bool = True,
) -> tuple[str | None, list[dict[str, Any]], dict[str, Any] | None]:
    """
    Make a streaming call to the model, optionally printing tokens as they arrive.

    Returns:
        Tuple of (text_content, tool_calls, assistant_message)
    """
    text_parts: list[str] = []
    tool_calls_by_index: dict[int, dict[str, Any]] = {}

    stream = client.chat.completions.create(
        model=model_name,
        messages=messages,
        tools=tools,
        temperature=0.7,
        stream=True,
    )

    for chunk in stream:
        if not chunk.choices:
            continue

        delta = chunk.choices[0].delta

        if delta.content:
            if print_output:
                print(delta.content, end="", flush=True)
            text_parts.append(delta.content)

        if delta.tool_calls:
            for tc in delta.tool_calls:
                idx = tc.index
                if idx not in tool_calls_by_index:
                    tool_calls_by_index[idx] = {
                        "id": tc.id or "",
                        "type": "function",
                        "function": {"name": "", "arguments": ""},
                    }
                if tc.id:
                    tool_calls_by_index[idx]["id"] = tc.id
                if tc.function:
                    if tc.function.name:
                        tool_calls_by_index[idx]["function"]["name"] = tc.function.name
                    if tc.function.arguments:
                        tool_calls_by_index[idx]["function"][
                            "arguments"
                        ] += tc.function.arguments

    if print_output and text_parts:
        print()

    text_content = "".join(text_parts) if text_parts else None
    tool_calls = [tool_calls_by_index[i] for i in sorted(tool_calls_by_index.keys())]

    if not text_content and not tool_calls:
        return None, [], None

    assistant_message: dict[str, Any] = {"role": "assistant"}
    if text_content:
        assistant_message["content"] = text_content
    if tool_calls:
        assistant_message["tool_calls"] = tool_calls

    return text_content, tool_calls, assistant_message


def _non_stream_response(
    client: OpenAI,
    model_name: str,
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]],
) -> tuple[str | None, list[dict[str, Any]], dict[str, Any] | None]:
    """Make a non-streaming call to the model."""
    response = client.chat.completions.create(
        model=model_name,
        messages=messages,
        tools=tools,
        temperature=0.7,
    )

    if not response.choices:
        return None, [], None

    message = response.choices[0].message
    text_content = message.content
    tool_calls = []

    if message.tool_calls:
        for tc in message.tool_calls:
            tool_calls.append(
                {
                    "id": tc.id,
                    "type": "function",
                    "function": {
                        "name": tc.function.name,
                        "arguments": tc.function.arguments,
                    },
                }
            )

    assistant_message: dict[str, Any] = {"role": "assistant"}
    if text_content:
        assistant_message["content"] = text_content
    if tool_calls:
        assistant_message["tool_calls"] = tool_calls

    return text_content, tool_calls, assistant_message


def run_agent_loop(
    env: CrmAgentEnv,
    client: OpenAI,
    model_name: str,
    tools: list[dict[str, Any]],
    messages: list[dict[str, Any]],
    max_iterations: int = 20,
    verbose: bool = False,
    stream: bool = False,
) -> tuple[str | None, list[dict[str, Any]]]:
    """
    Run the agent loop until it calls submit_answer or reaches max iterations.

    The agent can make multiple tool calls, getting results back each time,
    until it's ready to submit a final answer.

    Returns:
        Tuple of (final_answer or None, updated messages)
    """
    iteration = 0

    while iteration < max_iterations:
        iteration += 1

        try:
            if stream:
                text_content, tool_calls, assistant_message = _stream_response(
                    client=client,
                    model_name=model_name,
                    messages=messages,
                    tools=tools,
                    print_output=False,
                )
            else:
                text_content, tool_calls, assistant_message = _non_stream_response(
                    client=client,
                    model_name=model_name,
                    messages=messages,
                    tools=tools,
                )
        except Exception as e:
            print_colored(f"API Error: {e}", "red")
            return None, messages

        if assistant_message is None:
            print_colored("Error: No response from model", "red")
            return None, messages

        messages.append(assistant_message)

        if not tool_calls:
            return text_content, messages

        if text_content:
            print_colored("\nThinking:", "dim")
            print(text_content)

        final_answer = None

        for tc in tool_calls:
            tool_name = tc["function"]["name"]
            try:
                tool_args = json.loads(tc["function"]["arguments"])
            except json.JSONDecodeError:
                tool_args = {}

            if tool_name == "submit_answer":
                final_answer = tool_args.get("answer", "")
                print_colored("\n>>> Agent submitting answer", "cyan")
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tc["id"],
                        "content": "Answer submitted successfully.",
                    }
                )
                continue

            tool_display = format_tool_call_compact(tool_name, tool_args)
            print_colored(f"\n{tool_display}", "magenta")
            if verbose:
                print_colored(f"  Args: {json.dumps(tool_args, indent=2)}", "dim")

            parsed_tool_call = {"name": tool_name, "arguments": tool_args}
            parsed_action = env.parse_tool_call(parsed_tool_call)

            if not parsed_action.is_valid:
                obs_text = f"Error: {parsed_action.error_message}"
                print_colored(f"  Invalid: {parsed_action.error_message}", "red")
            else:
                try:
                    result = env.step(parsed_action.action)
                    obs_text = env.format_observation(result.observation)

                    display_text = (
                        obs_text[:500] + "..." if len(obs_text) > 500 else obs_text
                    )
                    print_colored(f"  Result: {display_text}", "green")
                except Exception as e:
                    obs_text = f"Error executing action: {e}"
                    print_colored(f"  {obs_text}", "red")

            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tc["id"],
                    "content": obs_text,
                }
            )

        if final_answer is not None:
            return final_answer, messages

    print_colored(f"\nMax iterations ({max_iterations}) reached", "yellow")
    return None, messages


def run_chat(
    env: CrmAgentEnv,
    client: OpenAI,
    model_name: str,
    tools: list[dict[str, Any]],
    system_prompt: str,
    initial_question: str | None = None,
    max_iterations: int = 20,
    verbose: bool = False,
    stream: bool = False,
) -> None:
    """
    Run an interactive chat session with the environment.

    The flow is:
    1. User asks a question
    2. Agent uses tools as needed (multiple calls in a loop)
    3. Agent calls submit_answer when ready
    4. Answer is displayed to user
    5. User can ask another question (loop back to 1)
    """
    print_colored("\nResetting environment...", "dim")
    env.reset()
    print_colored("Environment ready.\n", "green")

    messages: list[dict[str, Any]] = [{"role": "system", "content": system_prompt}]

    while True:
        if initial_question:
            question = initial_question
            initial_question = None
        else:
            print_colored("\nYou: ", "cyan")
            try:
                question = input().strip()
            except EOFError:
                break

        if not question:
            continue
        if question.lower() in ("quit", "exit", "q"):
            print_colored("Goodbye!", "yellow")
            break

        print_separator("=")

        messages.append({"role": "user", "content": question})

        answer, messages = run_agent_loop(
            env=env,
            client=client,
            model_name=model_name,
            tools=tools,
            messages=messages,
            max_iterations=max_iterations,
            verbose=verbose,
            stream=stream,
        )

        print_separator("=")
        if answer:
            print_colored("\nAssistant:", "blue")
            print(answer)
        else:
            print_colored("\nNo answer provided.", "yellow")
        print()


def main() -> None:
    """Main entry point for the chat CLI."""
    _check_dependencies()

    parser = argparse.ArgumentParser(
        description="Interactive chat CLI for the CRM environment",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python chat.py
    python chat.py --stream
    python chat.py --question "How many companies are there?"
    python chat.py --model anthropic/claude-sonnet-4 --stream

    # Using OpenAI directly:
    python chat.py --llm-base-url https://api.openai.com/v1 --model gpt-4o
        """,
    )

    parser.add_argument(
        "--crm-base-url",
        type=str,
        default="https://env.crm.construct-labs.com",
        help="Base URL for the CRM environment (default: https://env.crm.construct-labs.com)",
    )
    parser.add_argument(
        "--llm-base-url",
        type=str,
        default="https://openrouter.ai/api/v1",
        help="Base URL for OpenAI-compatible API (default: OpenRouter)",
    )
    parser.add_argument(
        "--model",
        type=str,
        default="anthropic/claude-sonnet-4.5",
        help="Model to use (default: anthropic/claude-sonnet-4.5)",
    )
    parser.add_argument(
        "--api-key",
        type=str,
        default=None,
        help="API key (default: uses OPENAI_API_KEY env var)",
    )
    parser.add_argument(
        "--question",
        type=str,
        default=None,
        help="Initial question/task for the agent",
    )
    parser.add_argument(
        "--max-iterations",
        type=int,
        default=20,
        help="Maximum tool call iterations per question (default: 20)",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print verbose output",
    )
    parser.add_argument(
        "--stream",
        action="store_true",
        help="Stream tokens from the LLM as they are generated",
    )

    args = parser.parse_args()

    api_key = args.api_key or os.getenv("OPENAI_API_KEY")
    if not api_key:
        print_colored("Error: No API key provided.", "red")
        print_colored("Set OPENAI_API_KEY in .env file or use --api-key", "dim")
        sys.exit(1)

    client = OpenAI(api_key=api_key, base_url=args.llm_base_url)

    print_colored(f"Connecting to CRM environment at {args.crm_base_url}...", "dim")
    try:
        env = CrmAgentEnv(base_url=args.crm_base_url)
    except Exception as e:
        print_colored(f"Error connecting to environment: {e}", "red")
        sys.exit(1)

    system_prompt = """You are a helpful assistant with access to CRM tools.

Use the provided tools to help the user with their requests. You can make multiple tool calls to complete a task.

When you have gathered enough information or completed the task, call submit_answer with your response to the user."""

    try:
        with env:
            run_chat(
                env=env,
                client=client,
                model_name=args.model,
                tools=env.tools,
                system_prompt=system_prompt,
                initial_question=args.question,
                max_iterations=args.max_iterations,
                verbose=args.verbose,
                stream=args.stream,
            )
    except KeyboardInterrupt:
        print_colored("\n\nSession interrupted.", "yellow")


if __name__ == "__main__":
    main()
