"""Interactive REPL for onefirewall-ai."""

from __future__ import annotations

import os
import platform
from datetime import datetime
from pathlib import Path
from typing import Optional

from prompt_toolkit import PromptSession
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.formatted_text import HTML
from prompt_toolkit.history import FileHistory
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.styles import Style
from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.panel import Panel

from onefirewall_ai._onefirewall_compat import strip_onefirewall_xml_for_display
from onefirewall_ai.config import HOME_DIR, get_provider, init_home, load_access

console = Console(highlight=False)

ORANGE = "#d4772c"
BLUE = "#4a9eff"
GREEN = "#3ddc84"
DIM = "#666666"
WHITE = "#e8e8e8"
YELLOW = "#fbbf24"
RED = "#f87171"

PROMPT_STYLE = Style.from_dict(
    {
        "prompt.gt": f"bold {ORANGE}",
        "toolbar": f"bg:{DIM} {WHITE}",
        "toolbar.model": f"bold {BLUE}",
        "toolbar.session": f"{YELLOW}",
        "toolbar.cwd": f"{GREEN}",
        "toolbar.turns": f"{DIM}",
    }
)

SLASH_COMMANDS = {
    "/help": "Show this help",
    "/session": "Show current session ID",
    "/new": "Start a new session (clears history)",
    "/clear": "Clear the screen",
    "/exit": "Exit the REPL",
    "/quit": "Exit the REPL",
    "/agent": "Toggle agent mode (tools on or off)",
}


def _print_help() -> None:
    console.print()
    console.print(f"[bold {BLUE}]Slash commands[/]")
    for command, description in SLASH_COMMANDS.items():
        console.print(f"  [bold {ORANGE}]{command:<12}[/] [dim]{description}[/]")
    console.print()
    console.print("  [dim]Alt+Enter or end a line with \\ for multi-line input[/]")
    console.print("  [dim]Up and Down arrows navigate history[/]")
    console.print()


def _print_welcome(session_id: str, agent_mode: bool) -> None:
    mode_label = "Agent" if agent_mode else "Chat"
    console.print()
    console.print(
        Panel(
            f"  [bold {ORANGE}]OneFirewall AI[/]\n"
            f"  Mode: [bold {BLUE}]{mode_label}[/]   Session: [bold {YELLOW}]{session_id}[/]\n"
            "  Type [bold]/help[/] for commands, [bold]/exit[/] to quit",
            border_style=ORANGE,
            padding=(0, 1),
            expand=False,
        )
    )
    console.print()


def _make_toolbar(session_id: str, agent_mode: bool, turns: int, model: str) -> HTML:
    cwd = Path.cwd().name
    mode = "Agent" if agent_mode else "Chat"
    return HTML(
        f"<toolbar> {mode} - "
        f"<toolbar.model>{model}</toolbar.model> - "
        f"<toolbar.session>{session_id}</toolbar.session> - "
        f"<toolbar.cwd>{cwd}</toolbar.cwd> - "
        f"<toolbar.turns>turns: {turns}</toolbar.turns> "
        f"</toolbar>"
    )


def _render_assistant_header() -> None:
    console.print(f"\n[bold {ORANGE}]OneFirewall[/] [dim]------------------------------[/]")


def _render_user_marker() -> None:
    console.print(f"\n[bold {ORANGE}]You[/] [dim]---------------------------------------[/]")


def _render_tool_result(ok: bool, summary: str) -> None:
    marker = f"[bold {GREEN}]ok[/]" if ok else f"[bold {RED}]x[/]"
    console.print(f"  {marker} [dim]{summary}[/]")


def _result_summary(result: object) -> str:
    text = str(result)
    if len(text) > 100:
        lines = text.splitlines()
        return f"{lines[0][:80]}... ({len(lines)} lines)"
    return text[:100]


def _reset_messages(messages: list[dict], system_prompt: str) -> None:
    messages.clear()
    messages.append({"role": "system", "content": system_prompt})


def _save_session(session_id: str, messages: list[dict]) -> None:
    sessions_dir = HOME_DIR / "sessions"
    sessions_dir.mkdir(exist_ok=True)
    path = sessions_dir / f"{session_id}.md"

    lines = [f"# Session: {session_id}\n"]
    for message in messages:
        role = message.get("role", "")
        content = message.get("content", "")
        if role == "system":
            continue
        if role == "user":
            lines.append(f"## User\n{content}\n")
        elif role == "assistant":
            lines.append(f"## Assistant\n{content}\n")
        elif role == "tool":
            lines.append(f"## Tool Result\n```\n{content[:500]}\n```\n")

    path.write_text("\n".join(lines), encoding="utf-8")


def run_repl(
    *,
    initial_prompt: Optional[str] = None,
    session_id: Optional[str] = None,
    agent_mode: bool = True,
) -> None:
    init_home()
    creds = load_access()
    if not creds.get("api_key") or not creds.get("endpoint"):
        console.print(f"[bold {RED}]Not logged in.[/] Run [bold]onefirewall-ai login[/] first.")
        return

    from onefirewall_ai.agent.loop import (
        _Spinner,
        _build_system_prompt,
        _execute_tool_with_spinner,
        _recover_empty_agent_reply,
        _stream_gateway,
        _wrap_agent_mode_prompt,
    )
    from onefirewall_ai.agent.router import SmartRouter
    from onefirewall_ai.agent.tools import get_tool_schemas

    model = creds.get("model", "openai/gpt-4o")
    if not session_id:
        session_id = f"repl_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

    history_dir = HOME_DIR / "history"
    history_dir.mkdir(exist_ok=True)
    history_file = history_dir / "repl_history.txt"

    router = SmartRouter(strategy=creds.get("strategy", "balanced"))
    router.initialize(creds)

    system_prompt = _build_system_prompt()
    messages: list[dict] = [{"role": "system", "content": system_prompt}]
    turns = 0

    _print_welcome(session_id, agent_mode)

    kb = KeyBindings()

    @kb.add("escape", "enter")
    def _newline(event):
        event.current_buffer.insert_text("\n")

    prompt_session = PromptSession(
        history=FileHistory(str(history_file)),
        auto_suggest=AutoSuggestFromHistory(),
        style=PROMPT_STYLE,
        key_bindings=kb,
        mouse_support=False,
        wrap_lines=True,
        enable_history_search=True,
    )

    def _toolbar() -> HTML:
        return _make_toolbar(session_id, agent_mode, turns, model)

    def _stream_live(provider, spinner, enabled_tools):
        current_text = ""
        provider_name = get_provider()  # "onefirewall" or "custom"
        with Live(Markdown(""), refresh_per_second=10, console=console) as live:
            def on_token(token: str) -> None:
                nonlocal current_text
                current_text += token
                live.update(Markdown(strip_onefirewall_xml_for_display(current_text)))

            return _stream_gateway(
                endpoint=provider.endpoint,
                api_key=provider.api_key,
                model=creds.get("model", provider.big_model),
                messages=messages,
                tools=enabled_tools,
                spinner=spinner,
                token_callback=on_token,
                provider=provider_name,
            )

    def _do_turn(user_text: str) -> None:
        nonlocal turns, agent_mode

        command = user_text.strip().lower().split()[0] if user_text.strip() else ""
        if command in ("/exit", "/quit"):
            raise SystemExit(0)
        if command == "/help":
            _print_help()
            return
        if command == "/clear":
            os.system("cls" if platform.system() == "Windows" else "clear")
            _print_welcome(session_id, agent_mode)
            return
        if command == "/session":
            console.print(f"\n  Session: [bold {YELLOW}]{session_id}[/]\n")
            return
        if command == "/new":
            _reset_messages(messages, system_prompt)
            console.print("\n  [dim]Conversation cleared.[/]\n")
            return
        if command == "/agent":
            agent_mode = not agent_mode
            label = "Agent (tools enabled)" if agent_mode else "Chat (no tools)"
            console.print(f"\n  Switched to [bold]{label}[/]\n")
            return

        turns += 1
        _render_user_marker()

        final_user_text = _wrap_agent_mode_prompt(user_text) if agent_mode else user_text
        messages.append({"role": "user", "content": final_user_text})

        spinner = _Spinner("Thinking")
        spinner.start()
        max_turns_per_prompt = 30

        try:
            for loop_turn in range(max_turns_per_prompt):
                provider = router.get_best_provider()
                try:
                    if loop_turn == 0:
                        _render_assistant_header()

                    enabled_tools = get_tool_schemas() if agent_mode else []
                    response = _stream_live(provider, spinner, enabled_tools)
                    spinner.stop()
                    if response.content.strip():
                        console.print()
                except Exception as exc:
                    spinner.stop()
                    console.print(f"\n[bold {RED}]API error:[/] {exc}\n")
                    return

                tool_calls = response.tool_calls if agent_mode else []
                content = _recover_empty_agent_reply(user_text, response.content or "", tool_calls)
                if content.strip() and not (response.content or "").strip():
                    console.print(strip_onefirewall_xml_for_display(content))
                    console.print()
                messages.append(
                    {
                        "role": "assistant",
                        "content": content,
                        **({"tool_calls": tool_calls} if tool_calls else {}),
                    }
                )

                if not tool_calls:
                    break

                tool_results = []
                for tool_call in tool_calls:
                    function = tool_call.get("function", {})
                    tool_name = function.get("name", "unknown")
                    raw_args = function.get("arguments", "{}")
                    try:
                        args = __import__("json").loads(raw_args)
                    except Exception:
                        args = {}

                    result = _execute_tool_with_spinner(tool_name, args)
                    ok = not result.startswith(("[ERROR]", "[DENIED]", "[ABORTED]"))
                    _render_tool_result(ok, _result_summary(result))
                    tool_results.append(
                        {
                            "role": "tool",
                            "tool_call_id": tool_call.get("id", ""),
                            "content": str(result),
                        }
                    )

                messages.extend(tool_results)
                spinner.set_message("Thinking")
                spinner.start()
        finally:
            spinner.stop()
            _save_session(session_id, messages)

    if initial_prompt:
        try:
            _do_turn(initial_prompt)
        except SystemExit:
            return
        except KeyboardInterrupt:
            console.print("\n[dim]Interrupted.[/]\n")

    while True:
        try:
            user_input = prompt_session.prompt(
                HTML('<prompt.gt>> </prompt.gt>'),
                style=PROMPT_STYLE,
                bottom_toolbar=_toolbar,
                multiline=False,
                rprompt=None,
            ).strip()
        except KeyboardInterrupt:
            console.print("\n[dim]Use /exit or Ctrl+D to quit.[/]\n")
            continue
        except EOFError:
            console.print("\n[dim]Goodbye.[/]\n")
            break

        if not user_input:
            continue

        while user_input.endswith("\\"):
            user_input = user_input[:-1].rstrip()
            try:
                more = prompt_session.prompt(
                    HTML('<prompt.gt>  </prompt.gt>'),
                    style=PROMPT_STYLE,
                ).strip()
                user_input = user_input + "\n" + more
            except (KeyboardInterrupt, EOFError):
                break

        try:
            _do_turn(user_input)
        except SystemExit:
            console.print("\n[dim]Goodbye.[/]\n")
            break
        except KeyboardInterrupt:
            console.print("\n[dim]Interrupted. Continue or /exit.[/]\n")
