from __future__ import annotations

import sys
import asyncio
import datetime
import shutil
import subprocess
from pathlib import Path

from rich.console import Console
from prompt_toolkit.formatted_text import FormattedText

from ..agent import stream_run
from ..engine.factory import get_engine
from ..engine.models import Model, compute_cost
from ..engine.types import Cost, Message, Role, Usage
from ..session import Session, append_messages, update_session_meta
from ..tools import TOOLS
from .display import show_banner, show_help, show_session, show_cost, show_tool_calls
from .input import InputSession, ReplCompleter

_console = Console()

_SYSTEM_PROMPT = (
    "You are Sun, an AI agent integrated into the SUNWÆE CLI. "
    f"Today is {datetime.datetime.now(datetime.timezone.utc).strftime('%Y-%m-%d')}. "
    "You help users manage their workspaces, notes, and tasks."
)


async def run(
    session: Session,
    history: list[Message],
    engine,
    workspace: str,
    model_name: str,
    provider: str,
    models_cache: list[Model],
) -> None:
    total_cost = Cost()
    total_usage = Usage()
    model_names = [m.name for m in models_cache]

    show_banner()

    _cwd = Path.cwd()
    try:
        cwd = "~/" + str(_cwd.relative_to(Path.home()))
    except ValueError:
        cwd = str(_cwd)
    total_cost_val = 0.0

    try:
        _branch = subprocess.run(
            ["git", "branch", "--show-current"],
            capture_output=True,
            text=True,
            cwd=_cwd,
        ).stdout.strip()
    except FileNotFoundError:
        _branch = ""

    def _get_status() -> FormattedText:
        width = shutil.get_terminal_size().columns
        branch_part = f"  │    {_branch}" if _branch else ""
        left = f"    {cwd}{branch_part}  │  󰚩   {provider}/{model_name} "
        right = f" 󰭹   {session.id[:8]}  │    {total_cost_val:.4f} "
        pad = max(1, width - len(left) - len(right))
        return FormattedText([("fg:#787d82", left + " " * pad + right)])

    _input = InputSession(ReplCompleter(model_names))

    while True:
        try:
            user_input = await _input.prompt_async(_get_status)
        except (EOFError, KeyboardInterrupt):
            break

        user_input = user_input.strip()
        if not user_input:
            continue

        sys.stdout.write("\n")
        sys.stdout.flush()

        # ── vim-style commands ────────────────────────────────────────────────
        if user_input.startswith(":"):
            parts = user_input[1:].split(maxsplit=1)
            cmd = parts[0].lower() if parts else ""
            arg = parts[1].strip() if len(parts) > 1 else ""

            if cmd in ("h", "help"):
                show_help()
            elif cmd in ("q", "quit", "exit"):
                break
            elif cmd == "clear":
                history.clear()
                _console.clear()
                continue
            elif cmd == "session":
                show_session(session, total_cost, total_usage)
            elif cmd == "cost":
                show_cost(total_cost, total_usage)
            elif cmd == "model":
                if arg:
                    found = next((m for m in models_cache if m.name == arg), None)
                    if not found:
                        _console.print(f"[red]Model {arg!r} not found.[/red]")
                    else:
                        try:
                            engine = get_engine(found.provider, found.name)
                            model_name = found.name
                            provider = found.provider
                        except ValueError as exc:
                            _console.print(f"[red]{exc}[/red]")
                else:
                    _console.print(f"  [bold]{model_name}[/bold]  {provider}")
            else:
                _console.print(f"[dim]Unknown command :{cmd!r}[/dim]")
            sys.stdout.write("\n")
            sys.stdout.flush()
            continue

        # ── chat turn ─────────────────────────────────────────────────────────
        user_msg = Message(role=Role.USER, content=user_input)

        if not session.title:
            session.title = user_input[:80]
            update_session_meta(workspace, session)

        all_msgs = (
            [Message(role=Role.SYSTEM, content=_SYSTEM_PROMPT)] + history + [user_msg]
        )
        new_messages: list[Message] = []
        reasoning_started = False
        content_started = False

        try:
            async for chunk in stream_run(
                all_msgs, TOOLS, engine, new_messages=new_messages
            ):
                if chunk.synthetic:
                    if chunk.tool_calls:
                        sys.stdout.write("\n")
                        sys.stdout.flush()
                        show_tool_calls(chunk.tool_calls)
                        reasoning_started = False
                        content_started = False
                    continue
                if chunk.reasoning_content:
                    if not reasoning_started:
                        reasoning_started = True
                    _console.print(
                        chunk.reasoning_content,
                        end="",
                        markup=False,
                        highlight=False,
                        style="italic #787d82",
                    )
                if chunk.content:
                    if not content_started:
                        sys.stdout.write("\n\n" if reasoning_started else "\n")
                        sys.stdout.flush()
                        content_started = True
                    _console.print(chunk.content, end="", markup=False, highlight=False)
                if chunk.stop_reason is not None:
                    if chunk.usage:
                        current_model_info = next(
                            (m for m in models_cache if m.name == model_name), None
                        )
                        if current_model_info:
                            chunk.cost = compute_cost(chunk.usage, current_model_info)
                    if chunk.cost:
                        total_cost.input += chunk.cost.input
                        total_cost.output += chunk.cost.output
                        total_cost.cache_read += chunk.cost.cache_read
                        total_cost.cache_write += chunk.cost.cache_write
                        total_cost.total += chunk.cost.total
                        total_cost_val = total_cost.total
                    if chunk.usage:
                        total_usage.input_tokens += chunk.usage.input_tokens
                        total_usage.output_tokens += chunk.usage.output_tokens
                        total_usage.total_tokens += chunk.usage.total_tokens
                        total_usage.cache_read_tokens += chunk.usage.cache_read_tokens
                        total_usage.cache_write_tokens += chunk.usage.cache_write_tokens
        except (KeyboardInterrupt, asyncio.CancelledError):
            sys.stdout.write("\n")
            _console.print("[dim]  interrupted[/dim]")

        sys.stdout.write("\n\n")
        sys.stdout.flush()

        append_messages(workspace, session.id, [user_msg] + new_messages)
        history.extend([user_msg] + new_messages)
