"""
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 google-genai python-dotenv

Usage:
    export CRM_AGENT_API_KEY=your-crm-api-key
    export GOOGLE_API_KEY=your-google-api-key
    python chat.py --help
    python chat.py
    python chat.py --question "How many companies are there?"
    python chat.py --model gemini-2.0-flash --stream
"""

from __future__ import annotations

import argparse
import json
import os
import sys

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

try:
    from google import genai
    from google.genai import types
except ImportError:
    genai = None  # type: ignore[assignment]
    types = None  # type: ignore[assignment]

from construct_labs_crm_env import CrmAgentEnv


def _check_dependencies() -> None:
    """Check that optional dependencies are installed."""
    missing = []
    if genai is None:
        missing.append("google-genai")
    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 convert_openai_tools_to_gemini(openai_tools: list[dict]) -> list[types.Tool]:
    """Convert OpenAI-format tools to Gemini format."""
    function_declarations = []

    for tool in openai_tools:
        if tool.get("type") != "function":
            continue

        func = tool["function"]
        name = func["name"]
        description = func.get("description", "")
        parameters = func.get("parameters", {})

        function_declarations.append(
            types.FunctionDeclaration(
                name=name,
                description=description,
                parameters=parameters if parameters else None,
            )
        )

    return [types.Tool(function_declarations=function_declarations)]


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: genai.Client,
    model_name: str,
    contents: list[types.Content],
    config: types.GenerateContentConfig,
    print_output: bool = True,
) -> tuple[list[str], list[types.FunctionCall], types.Content | None]:
    """
    Make a streaming call to the model, optionally printing tokens as they arrive.

    Returns:
        Tuple of (text_parts, function_calls, final_content)
    """
    text_parts: list[str] = []
    function_calls: list[types.FunctionCall] = []
    all_parts: list[types.Part] = []

    for chunk in client.models.generate_content_stream(
        model=model_name,
        contents=contents,
        config=config,
    ):
        if not chunk.candidates:
            continue

        candidate = chunk.candidates[0]

        for part in candidate.content.parts:
            if part.text:
                if print_output:
                    print(part.text, end="", flush=True)
                text_parts.append(part.text)
                all_parts.append(part)

            if part.function_call:
                function_calls.append(part.function_call)
                all_parts.append(part)

    if print_output and text_parts:
        print()

    if not all_parts:
        return [], [], None

    final_content = types.Content(role="model", parts=all_parts)
    full_text = "".join(text_parts)
    text_result = [full_text] if full_text else []

    return text_result, function_calls, final_content


def _non_stream_response(
    client: genai.Client,
    model_name: str,
    contents: list[types.Content],
    config: types.GenerateContentConfig,
) -> tuple[list[str], list[types.FunctionCall], types.Content | None]:
    """Make a non-streaming call to the model."""
    response = client.models.generate_content(
        model=model_name,
        contents=contents,
        config=config,
    )

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

    candidate = response.candidates[0]
    text_parts: list[str] = []
    function_calls: list[types.FunctionCall] = []

    for part in candidate.content.parts:
        if part.text:
            text_parts.append(part.text)
        if part.function_call:
            function_calls.append(part.function_call)

    return text_parts, function_calls, candidate.content


def run_agent_loop(
    env: CrmAgentEnv,
    client: genai.Client,
    model_name: str,
    tools: list[types.Tool],
    system_instruction: str,
    contents: list[types.Content],
    max_iterations: int = 20,
    verbose: bool = False,
    stream: bool = False,
) -> tuple[str | None, list[types.Content]]:
    """
    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 contents)
    """
    iteration = 0

    while iteration < max_iterations:
        iteration += 1

        config = types.GenerateContentConfig(
            tools=tools,
            system_instruction=system_instruction,
            temperature=0.7,
        )

        try:
            if stream:
                text_parts, function_calls, final_content = _stream_response(
                    client=client,
                    model_name=model_name,
                    contents=contents,
                    config=config,
                    print_output=False,
                )
            else:
                text_parts, function_calls, final_content = _non_stream_response(
                    client=client,
                    model_name=model_name,
                    contents=contents,
                    config=config,
                )
        except Exception as e:
            print_colored(f"API Error: {e}", "red")
            return None, contents

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

        contents.append(final_content)

        if not function_calls:
            if text_parts:
                return "\n".join(text_parts), contents
            return None, contents

        if text_parts:
            print_colored("\nThinking:", "dim")
            print("\n".join(text_parts))

        function_response_parts = []
        final_answer = None

        for func_call in function_calls:
            tool_name = func_call.name
            tool_args = dict(func_call.args) if func_call.args else {}

            if tool_name == "submit_answer":
                final_answer = tool_args.get("answer", "")
                print_colored("\n>>> Agent submitting answer", "cyan")
                function_response_parts.append(
                    types.Part.from_function_response(
                        name=func_call.name,
                        response={"result": "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")

            function_response_parts.append(
                types.Part.from_function_response(
                    name=func_call.name,
                    response={"result": obs_text},
                )
            )

        contents.append(types.Content(role="user", parts=function_response_parts))

        if final_answer is not None:
            return final_answer, contents

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


def run_chat(
    env: CrmAgentEnv,
    client: genai.Client,
    model_name: str,
    tools: list[types.Tool],
    system_instruction: 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")

    contents: list[types.Content] = []

    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("=")

        contents.append(
            types.Content(role="user", parts=[types.Part.from_text(text=question)])
        )

        answer, contents = run_agent_loop(
            env=env,
            client=client,
            model_name=model_name,
            tools=tools,
            system_instruction=system_instruction,
            contents=contents,
            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()

    # load_dotenv already called at module level if available

    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 gemini-2.0-flash --stream
        """,
    )

    parser.add_argument(
        "--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(
        "--model",
        type=str,
        default="gemini-2.0-flash",
        help="Gemini model to use (default: gemini-2.0-flash)",
    )
    parser.add_argument(
        "--api-key",
        type=str,
        default=None,
        help="Google API key (default: uses GOOGLE_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("GOOGLE_API_KEY")
    if not api_key:
        print_colored("Error: No Google API key provided.", "red")
        print_colored("Set GOOGLE_API_KEY in .env file or use --api-key", "dim")
        sys.exit(1)

    client = genai.Client(api_key=api_key)

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

    gemini_tools = convert_openai_tools_to_gemini(env.tools)

    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=gemini_tools,
                system_instruction=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()
