commands.py

"""Command handlers, session state, and input completion for Nova."""

from __future__ import annotations

import datetime
from dataclasses import dataclass, field
import pyperclip
from pathlib import Path
from typing import Callable, Optional

import ollama
from tools import load_memory, MEMORY_PATH
from prompt_toolkit.completion import Completer, Completion
from prompt_toolkit.styles import Style
from rich.align import Align
from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.rule import Rule
from rich.spinner import Spinner
from rich.text import Text

CONTEXT_DIR = Path(__file__).parent / "context"

SLASH_COMMANDS = [
    "/help", "/model", "/profile", "/profiles", "/models",
    "/image", "/file", "/clear", "/status", "/summary", "/copy", "/memory", "/exit",
]

prompt_style = Style.from_dict({"prompt": "#ffd700 bold"})


@dataclass
class Session:
    console: Console
    config: dict
    model: str
    profile_name: str
    history: list[dict]
    max_history: int
    pending_image: Optional[Path] = None


def build_profile(config: dict, profile_name: str, console: Console) -> tuple[str, list[dict]]:
    """Return (model, initial_history) for the given profile."""
    profiles = config.get("profiles", {})
    if profile_name not in profiles:
        console.print(f"[red]Unknown profile '{profile_name}'. Falling back to 'default'.[/red]")
        profile_name = "default"

    profile = profiles[profile_name]
    model = profile["model"]

    now = datetime.datetime.now()
    base = config.get("settings", {}).get("base_system_prompt", "").strip()
    base = base.format(date=now.strftime("%Y-%m-%d"), year=now.year)
    specific = profile.get("system_prompt", "").strip()
    facts = load_memory()
    memory_block = ("## What you remember:\n" + "\n".join(f"- {f}" for f in facts)) if facts else ""
    system_prompt = "\n\n".join(filter(None, [base, specific, memory_block]))

    history: list[dict] = []
    if system_prompt:
        history.append({"role": "system", "content": system_prompt})

    return model, history


def trim_history(history: list[dict], max_messages: int) -> list[dict]:
    """Drop oldest non-system messages when the limit is exceeded."""
    system = [m for m in history if m["role"] == "system"]
    non_system = [m for m in history if m["role"] != "system"]
    while len(non_system) > max_messages:
        non_system.pop(0)
    return system + non_system


def cmd_help(session: Session, argument: str) -> None:
    rows = [
        ("/help", "show this message"),
        ("/model <name>", "switch Ollama model"),
        ("/profile <name>", "load a profile (clears history)"),
        ("/profiles", "list available profiles"),
        ("/models", "list locally installed Ollama models"),
        ("/image <path>", "attach an image to the next message"),
        ("/file <path>", "inject a file's contents into context"),
        ("/clear", "clear conversation history"),
        ("/status", "show current model and profile"),
        ("/summary", "summarize the conversation so far"),
        ("/copy",    "copy the last response to clipboard"),
        ("/memory",  "show persistent memory"),
        ("/exit", "quit"),
    ]
    text = Text("\n")
    for cmd, desc in rows:
        text.append(f"  {cmd:<24}", style="color(220) bold")
        text.append(f"{desc}\n", style="color(245)")
    text.append("\n")
    session.console.print(text)


def cmd_clear(session: Session, argument: str) -> None:
    session.history = [m for m in session.history if m["role"] == "system"]
    session.console.print("[color(240)]history cleared[/color(240)]")


def cmd_status(session: Session, argument: str) -> None:
    non_sys = [m for m in session.history if m["role"] != "system"]
    text = Text("\n")
    text.append("  model    ", style="color(240)")
    text.append(session.model + "\n", style="color(220)")
    text.append("  profile  ", style="color(240)")
    text.append(session.profile_name + "\n", style="color(220)")
    text.append("  history  ", style="color(240)")
    text.append(f"{len(non_sys)} messages\n\n", style="color(220)")
    session.console.print(text)


def print_switch(session: Session, what: str) -> None:
    session.console.print()
    session.console.print(Rule(f"[color(166)]switched {what}[/color(166)]", style="color(240)", characters="-"))
    info = Text(justify="center")
    info.append("model  ", style="color(240)")
    info.append(session.model, style="color(214)")
    info.append("   profile  ", style="color(240)")
    info.append(session.profile_name, style="color(220)")
    session.console.print(Align(info, align="center"))
    session.console.print()


def cmd_model(session: Session, argument: str) -> None:
    if not argument:
        session.console.print("[red]Usage: /model <name>[/red]")
        return
    session.model = argument
    print_switch(session, "model")


def cmd_profile(session: Session, argument: str) -> None:
    if not argument:
        session.console.print("[red]Usage: /profile <name>[/red]")
        return
    session.profile_name = argument
    session.model, session.history = build_profile(session.config, argument, session.console)
    session.pending_image = None
    print_switch(session, "profile")


def cmd_profiles(session: Session, argument: str) -> None:
    profiles = session.config.get("profiles", {})
    text = Text("\n")
    for name in profiles:
        active = name == session.profile_name
        suffix = "  (active)" if active else ""
        style = "color(220) bold" if active else "color(245)"
        text.append(f"  {name}{suffix}\n", style=style)
    text.append("\n")
    session.console.print(text)


def cmd_models(session: Session, argument: str) -> None:
    try:
        result = ollama.list()
        models = result.get("models", [])
        if not models:
            session.console.print("[color(240)]No local models found.[/color(240)]")
            return
        text = Text("\n")
        for m in models:
            name = m.get("name") or m.get("model", "unknown")
            size = m.get("size", 0)
            size_str = f"{size / 1e9:.1f} GB" if size else ""
            text.append(f"  {name}", style="color(220)")
            if size_str:
                text.append(f"  {size_str}", style="color(240)")
            text.append("\n")
        text.append("\n")
        session.console.print(text)
    except Exception as e:
        session.console.print(f"[red]Could not fetch models: {e}[/red]")


def cmd_image(session: Session, argument: str) -> None:
    if not argument:
        session.console.print("[red]Usage: /image <path>[/red]")
        return
    image_path = Path(argument).expanduser()
    if not image_path.exists():
        session.console.print(f"[red]File not found: {image_path}[/red]")
        return
    session.pending_image = image_path
    session.console.print(f"[color(220)]image queued: {image_path.name}[/color(220)]")
    session.console.print("[color(240)]will be attached to your next message[/color(240)]")


def cmd_file(session: Session, argument: str) -> None:
    if not argument:
        session.console.print("[red]Usage: /file <path>[/red]")
        return
    candidate = CONTEXT_DIR / argument
    file_path = candidate if candidate.exists() else Path(argument).expanduser()
    if not file_path.exists():
        session.console.print(f"[red]File not found: {file_path}[/red]")
        return
    try:
        content = file_path.read_text(encoding="utf-8")
    except Exception as e:
        session.console.print(f"[red]Could not read file: {e}[/red]")
        return
    injection = f"Contents of `{file_path.name}`:\n\n```\n{content}\n```"
    session.history.append({"role": "user", "content": injection})
    session.history.append({"role": "assistant", "content": f"Got it. I've read `{file_path.name}` and will reference it."})
    session.history = trim_history(session.history, session.max_history)
    session.console.print(f"[color(220)]injected: {file_path.name}[/color(220)]")


def cmd_summary(session: Session, argument: str) -> None:
    non_sys = [m for m in session.history if m["role"] not in ("system", "tool")]
    if not non_sys:
        session.console.print("[color(240)]nothing to summarize[/color(240)]")
        return
    messages = [
        *[m for m in session.history if m["role"] == "system"],
        *non_sys,
        {"role": "user", "content": "Summarize our conversation so far in a few concise bullet points."},
    ]
    try:
        with Live(Spinner("dots", text=" Summarizing...", style="color(208)"), console=session.console, refresh_per_second=15):
            response = ollama.chat(model=session.model, messages=messages)
        session.console.print(Markdown(response.message.content or ""))
    except Exception as e:
        session.console.print(f"[red]{e}[/red]")


def cmd_memory(session: Session, argument: str) -> None:
    facts = load_memory()
    if not facts:
        session.console.print("[color(240)]nothing remembered yet[/color(240)]")
        return
    text = Text("\n")
    for i, fact in enumerate(facts, 1):
        text.append(f"  {i}. ", style="color(240)")
        text.append(f"{fact}\n", style="color(220)")
    text.append("\n")
    session.console.print(text)


def cmd_copy(session: Session, argument: str) -> None:
    last = next((m["content"] for m in reversed(session.history) if m["role"] == "assistant"), None)
    if not last:
        session.console.print("[color(240)]nothing to copy[/color(240)]")
        return
    try:
        pyperclip.copy(last)
        session.console.print("[color(240)]copied[/color(240)]")
    except Exception as e:
        session.console.print(f"[red]clipboard error: {e}[/red]")


REGISTRY: dict[str, Callable[[Session, str], None]] = {
    "/help":     cmd_help,
    "/clear":    cmd_clear,
    "/status":   cmd_status,
    "/model":    cmd_model,
    "/profile":  cmd_profile,
    "/profiles": cmd_profiles,
    "/models":   cmd_models,
    "/image":    cmd_image,
    "/file":     cmd_file,
    "/summary":  cmd_summary,
    "/copy":     cmd_copy,
    "/memory":   cmd_memory,
}


class SlashCompleter(Completer):
    def __init__(self, config: dict) -> None:
        self._config = config
        self._cached_models: list[str] = []

    def _profile_names(self) -> list[str]:
        return list(self._config.get("profiles", {}).keys())

    def _model_names(self) -> list[str]:
        if not self._cached_models:
            try:
                result = ollama.list()
                self._cached_models = [
                    m.get("name") or m.get("model", "")
                    for m in result.get("models", [])
                ]
            except Exception:
                pass
        return self._cached_models

    def get_completions(self, document, complete_event):
        text = document.text_before_cursor

        if not text.startswith("/"):
            return

        parts = text.split(" ", 1)
        command = parts[0].lower()

        if len(parts) == 1:
            for cmd in SLASH_COMMANDS:
                if cmd.startswith(command):
                    yield Completion(cmd[len(command):], display=cmd)
            return

        arg = parts[1]

        if command == "/profile":
            for name in self._profile_names():
                if name.startswith(arg):
                    yield Completion(name[len(arg):], display=name)

        elif command == "/model":
            for name in self._model_names():
                if name.startswith(arg):
                    yield Completion(name[len(arg):], display=name)