#!/usr/bin/env python3
"""
SAL-v4 Harness Full TUI
Terminal User Interface for interactive agent operation.
Uses Rich for rendering with proper input handling.
"""

import os
import sys
import time
import threading
from datetime import datetime
from typing import Optional, List, Dict
from queue import Queue, Empty
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

ENV_FILE = Path(__file__).parent / ".env"
if ENV_FILE.exists():
    with open(ENV_FILE) as f:
        for line in f:
            line = line.strip()
            if line and "=" in line and not line.startswith("#"):
                key, val = line.split("=", 1)
                os.environ.setdefault(key, val)

from rich.console import Console, Group
from rich.panel import Panel
from rich.text import Text
from rich.table import Table
from rich.style import Style
from rich.prompt import Prompt
from rich.markdown import Markdown

from core.zai_client import ZAiClient
from core.salience import SalienceEngine
from core.blackboard import Blackboard
from core.scheduler import Scheduler
from core.model_pool import MODEL_POOL, SLOT_SUMMARY, ThinkingMode
from agents.router import Router
from agents.tlp import TLPGateLoop, Phase


console = Console()


class AgentSession:
    """Manages agent state and communication."""

    def __init__(self, api_key: str):
        self.api_key = api_key
        self.bb = Blackboard()
        self.engine = SalienceEngine()
        self.gov = Scheduler(self.bb, self.engine)
        self.router = Router(self.bb, self.gov)
        self.tlp = TLPGateLoop(self.bb)

        self.model = "glm-4.7-flash"
        self.thinking_enabled = True
        self.preserve_thinking = True
        self.show_thinking = True

        self.messages: List[Dict] = []
        self.current_salience = 0.0
        self.current_phase = Phase.RECEIVE

        self.tokens_used = 0
        self.cost = 0.0
        self.errors = 0

    def chat(self, message: str) -> tuple:
        """Send message and get response. Returns (content, reasoning, error)."""

        thinking_config = None
        if self.thinking_enabled:
            thinking_config = {
                "type": "enabled",
                "clear_thinking": not self.preserve_thinking,
            }

        msgs = []
        for m in self.messages:
            if m["role"] == "user":
                msgs.append({"role": "user", "content": m["content"]})
            else:
                msg_dict = {"role": "assistant", "content": m["content"]}
                if self.preserve_thinking and m.get("reasoning"):
                    msg_dict["reasoning_content"] = m["reasoning"]
                msgs.append(msg_dict)
        msgs.append({"role": "user", "content": message})

        try:
            reasoning_parts = []
            content_parts = []

            with ZAiClient(self.api_key, timeout=60.0) as client:
                for r, c in client.chat_stream(
                    model=self.model,
                    messages=msgs,
                    thinking=thinking_config,
                    max_tokens=2048,
                ):
                    if r:
                        reasoning_parts.append(r)
                    if c:
                        content_parts.append(c)

            full_reasoning = "".join(reasoning_parts)
            full_content = "".join(content_parts)

            self.messages.append(
                {"role": "user", "content": message, "time": datetime.now()}
            )
            self.messages.append(
                {
                    "role": "assistant",
                    "content": full_content,
                    "reasoning": full_reasoning,
                    "time": datetime.now(),
                }
            )

            self.tokens_used += 100
            self.cost += 0.01

            salience = self.engine.calculate_l_sal(
                delta_a=0.5 if len(full_reasoning) > 100 else 0.2,
                r=0.1,
                m=0.3,
                c=0.8,
                phi=0.05,
                t=len(self.messages) * 0.1,
            )
            self.current_salience = salience

            self.tlp.step()
            self.current_phase = self.tlp.current_phase

            return full_content, full_reasoning, None

        except Exception as e:
            self.errors += 1
            return "", "", str(e)

    def clear(self):
        self.messages = []
        self.current_salience = 0.0
        self.tlp = TLPGateLoop(self.bb)
        self.current_phase = Phase.RECEIVE


def print_banner():
    console.clear()
    console.print(
        Panel.fit(
            "[bold cyan]SAL-v4 Harness TUI[/bold cyan]\n"
            "[dim]Salience-Optimized Agent Operating System[/dim]",
            border_style="cyan",
        )
    )


def print_status(session: AgentSession):
    """Print status bar."""
    console.print()

    table = Table(show_header=False, box=None, expand=True)
    table.add_column("key", style="dim", width=12)
    table.add_column("val", style="white", width=20)
    table.add_column("key2", style="dim", width=12)
    table.add_column("val2", style="white", width=20)

    thinking_str = "[green]ON[/green]" if session.thinking_enabled else "[red]OFF[/red]"
    preserve_str = (
        "[green]YES[/green]" if session.preserve_thinking else "[red]NO[/red]"
    )
    show_str = "[green]YES[/green]" if session.show_thinking else "[red]NO[/red]"

    table.add_row(
        "Model:", f"[yellow]{session.model}[/yellow]", "Thinking:", thinking_str
    )
    table.add_row("Preserve:", preserve_str, "Show:", show_str)
    table.add_row(
        "Salience:",
        f"[cyan]{session.current_salience:.2f}[/cyan]",
        "Phase:",
        f"[magenta]{session.current_phase.name}[/magenta]",
    )
    table.add_row("Tokens:", f"{session.tokens_used}", "Errors:", f"{session.errors}")

    console.print(Panel(table, title="Status", border_style="dim"))


def print_tlp_phases(session: AgentSession):
    """Print TLP phase progress."""
    phases = [
        ("RECEIVE", Phase.RECEIVE),
        ("WARN", Phase.WARN),
        ("PLAN", Phase.TENTATIVE_PLAN),
        ("MOVE", Phase.START_MOVEMENT),
        ("RECON", Phase.RECON),
        ("COMPLETE", Phase.COMPLETE_PLAN),
        ("ORDER", Phase.ISSUE_ORDER),
        ("REFINE", Phase.REFINE_SUPERVISE),
    ]

    phase_text = Text()
    for i, (name, p) in enumerate(phases):
        if p == session.current_phase:
            phase_text.append(f">>{name}", style="bold yellow")
        elif p.value < session.current_phase.value:
            phase_text.append(f"  {name}", style="dim green")
        else:
            phase_text.append(f"  {name}", style="dim")
        if i < len(phases) - 1:
            phase_text.append(" -> ")

    console.print(Panel(phase_text, title="TLP Phase", border_style="cyan"))


def print_history(session: AgentSession, limit: int = 5):
    """Print conversation history."""
    console.print()
    console.print("[bold blue]Recent Messages:[/bold blue]")

    for msg in session.messages[-limit * 2 :]:
        role = msg["role"]
        content = msg["content"]
        reasoning = msg.get("reasoning", "")
        time_str = (
            msg["time"].strftime("%H:%M:%S")
            if isinstance(msg["time"], datetime)
            else ""
        )

        if role == "user":
            console.print(f"\n[bold blue][{time_str}] You:[/bold blue]")
            console.print(f"  {content[:500]}{'...' if len(content) > 500 else ''}")
        else:
            if reasoning and session.show_thinking:
                console.print(f"\n[dim italic][{time_str}] Thinking:[/dim italic]")
                console.print(
                    f"  [dim]{reasoning[:300]}{'...' if len(reasoning) > 300 else ''}[/dim]"
                )
            console.print(f"\n[bold green][{time_str}] Agent:[/bold green]")
            console.print(f"  {content[:500]}{'...' if len(content) > 500 else ''}")


def print_help():
    """Print help message."""
    console.print(
        Panel(
            "[bold]Commands:[/bold]\n"
            "  /thinking    - Toggle thinking mode\n"
            "  /preserve    - Toggle preserved thinking\n"
            "  /show        - Toggle showing thinking output\n"
            "  /model NAME  - Switch model (glm-4.7-flash, glm-4.7, glm-5)\n"
            "  /clear       - Clear conversation\n"
            "  /status      - Show full status\n"
            "  /pool        - Show model pool\n"
            "  /help        - Show this help\n"
            "  /quit        - Exit",
            title="Help",
            border_style="yellow",
        )
    )


def print_model_pool():
    """Print model pool configuration."""
    console.print(Panel(SLOT_SUMMARY, title="Model Pool", border_style="green"))


def main():
    api_key = os.environ.get("ZAI_API_KEY")

    if not api_key:
        console.print("[red]ERROR: ZAI_API_KEY not set[/red]")
        console.print("Usage: ZAI_API_KEY=your-key python tui.py")
        sys.exit(1)

    session = AgentSession(api_key)

    print_banner()
    console.print("\n[dim]Type a message or /help for commands[/dim]")

    while True:
        try:
            print_status(session)

            user_input = console.input("\n[bold white]You>[/bold white] ").strip()

            if not user_input:
                continue

            if user_input.startswith("/"):
                cmd_parts = user_input[1:].split(maxsplit=1)
                cmd = cmd_parts[0].lower()
                args = cmd_parts[1] if len(cmd_parts) > 1 else ""

                if cmd in ("quit", "exit", "q"):
                    console.print("\n[yellow]Goodbye![/yellow]")
                    break

                elif cmd == "thinking":
                    session.thinking_enabled = not session.thinking_enabled
                    console.print(
                        f"[green]Thinking: {'ON' if session.thinking_enabled else 'OFF'}[/green]"
                    )

                elif cmd == "preserve":
                    session.preserve_thinking = not session.preserve_thinking
                    console.print(
                        f"[green]Preserve: {'YES' if session.preserve_thinking else 'NO'}[/green]"
                    )

                elif cmd == "show":
                    session.show_thinking = not session.show_thinking
                    console.print(
                        f"[green]Show thinking: {'YES' if session.show_thinking else 'NO'}[/green]"
                    )

                elif cmd == "model":
                    models = ["glm-4.7-flash", "glm-4.7-flashx", "glm-4.7", "glm-5"]
                    if args in models:
                        session.model = args
                        console.print(f"[green]Model: {args}[/green]")
                    else:
                        console.print(
                            f"[yellow]Available: {', '.join(models)}[/yellow]"
                        )

                elif cmd == "clear":
                    session.clear()
                    console.print("[green]Cleared[/green]")

                elif cmd == "status":
                    print_status(session)
                    print_tlp_phases(session)
                    print_history(session, limit=10)

                elif cmd == "pool":
                    print_model_pool()

                elif cmd == "help":
                    print_help()

                else:
                    console.print(f"[red]Unknown command: /{cmd}[/red]")

                continue

            console.print("\n[dim]Processing...[/dim]")

            content, reasoning, error = session.chat(user_input)

            if error:
                console.print(f"\n[red]Error: {error}[/red]")
            else:
                if reasoning and session.show_thinking:
                    console.print(f"\n[dim italic]Thinking:[/dim italic]")
                    console.print(
                        Panel(
                            reasoning[:500] + ("..." if len(reasoning) > 500 else ""),
                            border_style="dim",
                            style="dim",
                        )
                    )

                console.print(f"\n[bold green]Agent:[/bold green]")
                console.print(Panel(content, border_style="green"))

        except KeyboardInterrupt:
            console.print("\n[yellow]Interrupted - type /quit to exit[/yellow]")
        except Exception as e:
            console.print(f"\n[red]Error: {e}[/red]")


if __name__ == "__main__":
    main()