"""Agent loop for OneFirewall AI with branded prompting and XML tool parsing."""

from __future__ import annotations

import json
import platform
import re
import signal
import subprocess
import sys
import threading
import time
from datetime import datetime
from typing import Callable, Optional

import requests as _requests
from rich.console import Console
from rich.markup import escape
from rich.panel import Panel
from rich.rule import Rule

from onefirewall_ai.config import (
    MEMORY_FILE,
    SOUL_FILE,
    create_session_file,
    get_provider,
    get_session_file,
    is_logged_in,
    load_access,
)
from onefirewall_ai._onefirewall_compat import extract_onefirewall_xml_tool_calls
from onefirewall_ai._onefirewall_snapshot import (
    CURATED_STATIC_PROMPT_SECTIONS,
    DEFAULT_AGENT_PROMPT,
)
from onefirewall_ai.agent.router import SmartRouter
from onefirewall_ai.agent.tools import (
    execute_tool,
    get_session_todos,
    get_tool_schemas,
    reset_session_permissions,
    reset_session_todos,
)

console = Console(force_terminal=True)

MAX_TURNS = 50
MAX_RETRIES = 3
_abort_flag = threading.Event()


def _handle_sigint(sig, frame):
    _abort_flag.set()
    sys.stderr.write("\r\033[K")
    sys.stderr.flush()
    console.print("\n[bold yellow]Interrupted - stopping after the current tool.[/bold yellow]")


class _Spinner:
    """Terminal spinner that clears cleanly before streamed output."""

    def __init__(self, message: str):
        self._message = message
        self._status = None

    def set_message(self, message: str) -> None:
        self._message = message

    def start(self) -> None:
        if self._status is None:
            sys.stdout.write("\r\033[K")
            sys.stdout.flush()
            self._status = console.status(f"  {self._message}", spinner="line")
            self._status.start()

    def stop(self) -> None:
        if self._status is not None:
            self._status.stop()
            self._status = None
        sys.stdout.write("\r\033[K")
        sys.stdout.flush()

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, *_):
        self.stop()


class CostTracker:
    def __init__(self) -> None:
        self.input_tokens = 0
        self.output_tokens = 0
        self.turns = 0

    def record(self, usage: dict) -> None:
        self.input_tokens += usage.get("prompt_tokens", 0)
        self.output_tokens += usage.get("completion_tokens", 0)
        self.turns += 1

    def summary(self) -> str:
        total = self.input_tokens + self.output_tokens
        if total == 0:
            return f"Turns: {self.turns}"
        return (
            f"Tokens: {self.input_tokens:,} in + {self.output_tokens:,} out = {total:,} total"
            f"  |  Turns: {self.turns}"
        )


def _load_session_messages(session_id: str) -> list[dict]:
    path = get_session_file(session_id)
    messages: list[dict] = []
    if not path.exists():
        return messages

    try:
        text = path.read_text(encoding="utf-8")
    except Exception:
        return messages

    for block in text.split("---"):
        lines = block.strip().splitlines()
        user_msg = ""
        assistant_msg = ""
        for line in lines:
            if line.startswith("**User:** "):
                user_msg = line[len("**User:** "):]
            elif line.startswith("**Assistant:** "):
                assistant_msg = line[len("**Assistant:** "):]
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    return messages


def _save_turn(session_id: str, user_msg: str, assistant_msg: str) -> None:
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
    entry = (
        f"\n## [{timestamp}]\n"
        f"**User:** {user_msg}\n"
        f"**Assistant:** {assistant_msg}\n\n"
        f"---\n"
    )
    with open(get_session_file(session_id), "a", encoding="utf-8") as handle:
        handle.write(entry)


def _build_system_prompt() -> str:
    soul = SOUL_FILE.read_text(encoding="utf-8") if SOUL_FILE.exists() else ""
    memory = MEMORY_FILE.read_text(encoding="utf-8") if MEMORY_FILE.exists() else ""

    cwd = subprocess.os.getcwd()
    is_windows = platform.system() == "Windows"
    shell = subprocess.os.environ.get("SHELL", "powershell" if is_windows else "bash")

    try:
        subprocess.check_output(
            ["git", "rev-parse", "--git-dir"],
            cwd=cwd,
            stderr=subprocess.DEVNULL,
            timeout=2,
        )
        is_git = "Yes"
    except Exception:
        is_git = "No"

    environment_section = "\n".join(
        [
            "# Environment",
            f"- Operating System: {platform.system()} ({platform.machine()})",
            f"- Working Directory: {cwd}",
            f"- Is Git Repository: {is_git}",
            f"- Shell: {shell}",
            f"- Date: {datetime.now().strftime('%Y-%m-%d')}",
        ]
    )
    onefirewall_guidance = "\n".join(
        [
            "# OneFirewall guidance",
            "- Before destructive, hard-to-reverse, or externally visible actions, check with the user first.",
            "- Use todo_write proactively for multi-step work and keep it up to date.",
            "- Do not create markdown reports, summary files, or extra artifacts unless the user asks for them.",
        ]
    )
    working_with_user = "\n".join(
        [
            "# Working with the user",
            "- Before the first tool call for a task, send one short sentence telling the user what you are about to do.",
            "- Before editing files or running a substantial command, send a brief progress update.",
            "- During multi-step work, keep the user informed with short natural updates at meaningful milestones.",
            "- Keep progress updates concise and practical. Do not dump long reasoning.",
        ]
    )

    sections = [
        "You are OneFirewall AI, an autonomous software engineering agent for the OneFirewall AI Gateway. Help users with complex development tasks in their workspace.",
        f"# OneFirewall agent baseline\n{DEFAULT_AGENT_PROMPT}",
        environment_section,
        CURATED_STATIC_PROMPT_SECTIONS["system"],
        CURATED_STATIC_PROMPT_SECTIONS["doing_tasks"],
        CURATED_STATIC_PROMPT_SECTIONS["executing_actions_with_care"],
        CURATED_STATIC_PROMPT_SECTIONS["using_your_tools"],
        onefirewall_guidance,
        working_with_user,
        CURATED_STATIC_PROMPT_SECTIONS["tone_and_style"],
        CURATED_STATIC_PROMPT_SECTIONS["communicating_with_user"],
    ]

    if soul.strip():
        sections.append(soul.strip())
    if memory.strip():
        sections.append(memory.strip())

    return "\n\n".join(section for section in sections if section.strip())


def _wrap_agent_mode_prompt(user_text: str) -> str:
    return (
        f"{user_text}\n\n"
        "[IMPORTANT: You are in AGENT MODE.\n"
        "- Complete the task fully without gold-plating or leaving it half-done\n"
        "- If the user's message is conversational (for example: hey, hello, thanks), reply normally and do not force tool use\n"
        "- If the request requires inspecting files, running commands, or making changes, first send one short sentence telling the user what you are about to do\n"
        "- Then use your tools to take action directly when tools are needed\n"
        "- Keep progress updates brief, natural, and useful\n"
        "- Do NOT create extra report files or markdown summaries unless the user asks for them\n"
        "- Prefer replying with findings in the chat unless a file is explicitly requested\n"
        "- Use write_file to create all files; do NOT write code in your text response\n"
        "- Use bash to run commands; do NOT describe commands for the user to run\n"
        "- For long-running dev servers or watchers, use bash with run_in_background=true instead of appending & to the command\n"
        "- If the user asks you to run or start an app/server, launch it, report the URL or startup failure once, and stop there unless they explicitly ask you to debug or fix it\n"
        "- If a startup health check fails, summarize the failure instead of repeatedly killing and restarting the server unless the user asked for troubleshooting\n"
        "- Use edit_file to fix bugs; do NOT say 'change line X to Y'\n"
        "Start immediately after your short progress update, and use todo_write first if the task is multi-step.]"
    )


def _looks_like_casual_message(user_text: str) -> bool:
    normalized = re.sub(r"\s+", " ", user_text.strip().lower())
    if not normalized:
        return False
    return bool(
        re.match(
            r"^(hi|hey|hello|yo|sup|good (morning|afternoon|evening)|thanks|thank you|ok|okay|cool)\b",
            normalized,
        )
    )


def _recover_empty_agent_reply(user_text: str, content: str, tool_calls: list[dict]) -> str:
    if tool_calls or content.strip():
        return content

    normalized = re.sub(r"\s+", " ", user_text.strip().lower())
    if normalized in {"continue", "go on", "go ahead", "proceed"}:
        return "I hit an empty model response just then. Please send that again and I'll continue from the last step."
    if normalized.startswith(("thanks", "thank you")):
        return "You're welcome. What would you like to work on next?"
    if _looks_like_casual_message(user_text):
        return "Hey! What would you like to work on?"
    return "I didn't get a usable response from the model just then. Please retry your last message."


class GatewayResponse:
    def __init__(self, content: str, tool_calls: list[dict], finish_reason: str, usage: dict):
        self.content = content
        self.tool_calls = tool_calls
        self.finish_reason = finish_reason
        self.usage = usage


def _stream_gateway(
    endpoint: str,
    api_key: str,
    model: str,
    messages: list[dict],
    tools: list[dict],
    spinner: _Spinner,
    token_callback: Optional[Callable[[str], None]] = None,
    provider: str = "onefirewall",
) -> GatewayResponse:
    # OneFirewall uses /api/v1/…; other OpenAI-compatible providers use /v1/…
    path = "/api/v1/chat/completions" if provider == "onefirewall" else "/v1/chat/completions"
    url = endpoint.rstrip("/") + path
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "Accept": "text/event-stream, text/plain",
    }
    payload = {
        "model": model,
        "messages": messages,
        "stream": True,
    }
    # OneFirewall-specific payload key
    if provider == "onefirewall":
        payload["pii"] = "disabled"

    allow_tool_calls = bool(tools)
    if allow_tool_calls:
        payload["tools"] = tools
        payload["tool_choice"] = "auto"

    content_parts: list[str] = []
    tool_calls_acc: dict[int, dict] = {}
    finish_reason = "stop"
    usage: dict = {}
    spinner_stopped = False

    def stop_spinner_once() -> None:
        nonlocal spinner_stopped
        if not spinner_stopped:
            spinner.stop()
            spinner_stopped = True

    try:
        response = _requests.post(
            url,
            headers=headers,
            json=payload,
            stream=True,
            timeout=(10, 120),
        )
        if response.status_code != 200:
            stop_spinner_once()
            raise RuntimeError(f"HTTP {response.status_code}: {response.text[:500]}")

        is_sse = "event-stream" in response.headers.get("Content-Type", "")
        if is_sse:
            for raw_line in response.iter_lines():
                if _abort_flag.is_set():
                    break
                if not raw_line:
                    continue

                line = raw_line.decode("utf-8") if isinstance(raw_line, bytes) else raw_line
                if not line.startswith("data:"):
                    continue
                payload_str = line[5:].strip()
                if payload_str == "[DONE]":
                    break

                try:
                    data = json.loads(payload_str)
                except json.JSONDecodeError:
                    continue

                if "usage" in data:
                    usage = data["usage"]

                choices = data.get("choices") or []
                if not choices:
                    continue

                choice = choices[0]
                delta = choice.get("delta", {})
                if choice.get("finish_reason"):
                    finish_reason = choice["finish_reason"]

                token = delta.get("content") or ""
                if token:
                    stop_spinner_once()
                    if token_callback:
                        token_callback(token)
                    else:
                        sys.stdout.write(token)
                        sys.stdout.flush()
                    content_parts.append(token)

                if allow_tool_calls:
                    for tool_delta in delta.get("tool_calls", []):
                        index = tool_delta.get("index", 0)
                        if index not in tool_calls_acc:
                            tool_calls_acc[index] = {
                                "id": "",
                                "type": "function",
                                "function": {"name": "", "arguments": ""},
                            }
                        tool_call = tool_calls_acc[index]
                        if tool_delta.get("id"):
                            tool_call["id"] += tool_delta["id"]
                        function_delta = tool_delta.get("function", {})
                        if function_delta.get("name"):
                            tool_call["function"]["name"] += function_delta["name"]
                        if function_delta.get("arguments"):
                            tool_call["function"]["arguments"] += function_delta["arguments"]
        else:
            for chunk in response.iter_content(chunk_size=None):
                if _abort_flag.is_set():
                    break
                if not chunk:
                    continue
                text = chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk
                stop_spinner_once()
                if token_callback:
                    token_callback(text)
                else:
                    sys.stdout.write(text)
                    sys.stdout.flush()
                content_parts.append(text)
    except _requests.exceptions.ConnectionError:
        stop_spinner_once()
        raise RuntimeError(
            f"Cannot connect to '{endpoint}'.\nRun `onefirewall-ai login` to update your configuration."
        )
    except _requests.exceptions.Timeout:
        stop_spinner_once()
        raise RuntimeError("Request timed out (120s). Check your connection.")
    except _requests.exceptions.RequestException as exc:
        stop_spinner_once()
        raise RuntimeError(f"Request failed: {exc}")

    stop_spinner_once()

    content = "".join(content_parts)
    tool_calls = [tool_calls_acc[index] for index in sorted(tool_calls_acc.keys())]
    if allow_tool_calls and not tool_calls:
        content, xml_tool_calls = extract_onefirewall_xml_tool_calls(content)
        if xml_tool_calls:
            tool_calls = xml_tool_calls

    return GatewayResponse(
        content=content,
        tool_calls=tool_calls,
        finish_reason=finish_reason,
        usage=usage,
    )


_TOOL_ICONS = {
    "bash": "*",
    "read_file": "o",
    "write_file": "W",
    "edit_file": "E",
    "glob_files": "G",
    "grep_files": "S",
    "list_directory": "L",
    "web_fetch": "W",
    "ask_user": "?",
    "todo_write": "T",
}

_TOOL_COLORS = {
    "bash": "green",
    "write_file": "yellow",
    "edit_file": "yellow",
    "read_file": "blue",
    "glob_files": "blue",
    "grep_files": "blue",
    "list_directory": "blue",
    "web_fetch": "magenta",
    "ask_user": "cyan",
    "todo_write": "dim",
}


def _print_tool_header(name: str, args: dict) -> None:
    icon = _TOOL_ICONS.get(name, "*")
    color = _TOOL_COLORS.get(name, "white")
    arg_summary = ""
    if args:
        first_key = next(iter(args))
        first_value = str(args[first_key])
        if len(first_value) > 80:
            first_value = first_value[:80] + "..."
        arg_summary = f"({escape(first_value)})"

    display_name = name.replace("_", " ").title().replace(" ", "")
    console.print(f"\n  [{color}]{icon} {display_name}{arg_summary}[/{color}]")

    for key, value in list(args.items())[1:]:
        value_str = str(value)
        if len(value_str) > 100:
            value_str = value_str[:100] + "..."
        console.print(f"    [dim]{key}: {escape(value_str)}[/dim]")


def _print_tool_result(result: str, is_error: bool) -> None:
    preview = result[:400]
    if len(result) > 400:
        preview += f"\n  ... ({len(result) - 400:,} more chars)"
    if is_error:
        console.print(f"  [red]x {escape(preview)}[/red]")
    else:
        console.print(f"  [dim]{escape(preview)}[/dim]")


def _tool_progress_message(name: str, args: dict) -> str:
    def shorten(value: object, limit: int = 60) -> str:
        text = re.sub(r"\s+", " ", str(value or "")).strip()
        if len(text) > limit:
            return text[:limit] + "..."
        return text

    match name:
        case "bash":
            command = shorten(args.get("command"), limit=70)
            return f"Running bash: {command}" if command else "Running bash command..."
        case "read_file":
            return f"Reading {shorten(args.get('path'))}..." if args.get("path") else "Reading file..."
        case "write_file":
            return f"Writing {shorten(args.get('path'))}..." if args.get("path") else "Writing file..."
        case "edit_file":
            return f"Editing {shorten(args.get('path'))}..." if args.get("path") else "Editing file..."
        case "glob_files":
            return "Finding matching files..."
        case "grep_files":
            return "Searching files..."
        case "list_directory":
            return f"Listing {shorten(args.get('path', '.'))}..."
        case "web_fetch":
            return f"Fetching {shorten(args.get('url'))}..." if args.get("url") else "Fetching URL..."
        case _:
            return f"Running {name.replace('_', ' ')}..."


def _execute_tool_with_spinner(name: str, args: dict) -> str:
    tool_spinner: _Spinner | None = None

    def start_spinner() -> None:
        nonlocal tool_spinner
        if name in {"ask_user", "todo_write"}:
            return
        tool_spinner = _Spinner(_tool_progress_message(name, args))
        tool_spinner.start()

    try:
        return execute_tool(name, args, on_before_execute=start_spinner)
    finally:
        if tool_spinner is not None:
            tool_spinner.stop()


def run_agent(prompt: str, session_id: str | None = None) -> None:
    if not is_logged_in():
        console.print(
            Panel(
                "[bold yellow]Not logged in.[/bold yellow]\nRun [bold cyan]onefirewall-ai login[/bold cyan] first.",
                border_style="yellow",
            )
        )
        sys.exit(1)

    access = load_access()
    router = SmartRouter(strategy=access.get("strategy", "balanced"))
    router.initialize(access)
    provider_obj = router.get_best_provider()

    endpoint = provider_obj.endpoint
    api_key = provider_obj.api_key
    model = access.get("model", provider_obj.big_model)
    provider_name = get_provider()  # "onefirewall" or "custom"

    reset_session_permissions()
    reset_session_todos()
    _abort_flag.clear()

    old_handler = signal.getsignal(signal.SIGINT)
    signal.signal(signal.SIGINT, _handle_sigint)

    if not session_id:
        session_id = f"agent_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    create_session_file(session_id)

    system_prompt = _build_system_prompt()
    messages: list[dict] = []
    if system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt})
    messages.extend(_load_session_messages(session_id))
    messages.append({"role": "user", "content": _wrap_agent_mode_prompt(prompt)})

    tools = get_tool_schemas()
    cost = CostTracker()

    console.print()
    console.print(Rule("[bold cyan]OneFirewall Agent[/bold cyan]", style="cyan"))
    console.print(f"  [dim]Model:[/dim] [cyan]{model}[/cyan]  [dim]Session:[/dim] [dim]{session_id}[/dim]")
    console.print()

    turn = 0
    final_answer = ""

    try:
        while turn < MAX_TURNS and not _abort_flag.is_set():
            turn += 1
            spinner = _Spinner("Thinking..." if turn == 1 else f"Thinking... (turn {turn})")
            spinner.start()

            response: GatewayResponse | None = None
            last_error = ""

            for attempt in range(MAX_RETRIES):
                try:
                    response = _stream_gateway(
                        endpoint, api_key, model, messages, tools, spinner,
                        provider=provider_name,
                    )
                    break
                except RuntimeError as exc:
                    spinner.stop()
                    last_error = str(exc)
                    if attempt + 1 < MAX_RETRIES:
                        wait_seconds = 2 ** attempt
                        console.print(f"\n[yellow]Retrying in {wait_seconds}s... ({exc})[/yellow]")
                        time.sleep(wait_seconds)
                        spinner = _Spinner(f"Retrying... ({attempt + 2}/{MAX_RETRIES})")
                        spinner.start()

            if response is None:
                console.print(f"\n[bold red]x Failed:[/bold red] {last_error}")
                return

            cost.record(response.usage)
            content = _recover_empty_agent_reply(prompt, response.content, response.tool_calls)
            tool_calls = response.tool_calls

            if content and not response.content.strip():
                sys.stdout.write(content)
                sys.stdout.write("\n")
                sys.stdout.flush()
            elif content:
                sys.stdout.write("\n")
                sys.stdout.flush()

            assistant_message: dict = {"role": "assistant", "content": content or ""}
            if tool_calls:
                assistant_message["tool_calls"] = tool_calls
            messages.append(assistant_message)

            if not tool_calls:
                final_answer = content
                break

            tool_results: list[dict] = []
            for tool_call in tool_calls:
                if _abort_flag.is_set():
                    tool_results.append(
                        {
                            "role": "tool",
                            "tool_call_id": tool_call.get("id", ""),
                            "content": "[ABORTED] User interrupted.",
                        }
                    )
                    continue

                function = tool_call.get("function", {})
                tool_name = function.get("name", "unknown")
                raw_args = function.get("arguments", "{}")
                try:
                    args = json.loads(raw_args) if raw_args.strip() else {}
                except json.JSONDecodeError:
                    args = {}

                _print_tool_header(tool_name, args)
                result = _execute_tool_with_spinner(tool_name, args)
                is_error = result.startswith(("[ERROR]", "[DENIED]", "[ABORTED]"))
                _print_tool_result(result, is_error)
                tool_results.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call.get("id", ""),
                        "content": result,
                    }
                )

            messages.extend(tool_results)
    finally:
        signal.signal(signal.SIGINT, old_handler)

    console.print()
    if _abort_flag.is_set():
        console.print(
            Panel(
                f"[yellow]Interrupted after {turn} turn(s).[/yellow]\nResume: [cyan]onefirewall-ai agent --session {session_id} \"continue\"[/cyan]",
                border_style="yellow",
                title="Interrupted",
            )
        )
    elif turn >= MAX_TURNS and not final_answer:
        console.print(
            Panel(
                f"[yellow]Reached {MAX_TURNS} turns - task may be incomplete.[/yellow]\nResume: [cyan]onefirewall-ai agent --session {session_id} \"continue\"[/cyan]",
                border_style="yellow",
                title="Turn Limit",
            )
        )

    if final_answer and not final_answer.startswith("[ERROR]"):
        _save_turn(session_id, prompt, final_answer)

    console.print(Rule("[bold cyan]Done[/bold cyan]", style="cyan"))
    console.print(f"  [dim]{cost.summary()}[/dim]")
    console.print(f"  [dim]Session:[/dim] [dim]{get_session_file(session_id)}[/dim]")

    remaining = [todo for todo in get_session_todos() if todo.get("status") != "completed"]
    if remaining:
        console.print(f"\n  [yellow]! {len(remaining)} task(s) still pending[/yellow]")

    console.print()

