"""glint-router-1m: route prompts, edit the policy, teach it new categories.

  python infer.py "write a python function that reverses a linked list"
  python infer.py --file prompts.txt
  python infer.py --policy my_policy.json "summarize this thread"
  python infer.py --add-category refunds examples.txt
  python infer.py --demo

everything runs on cpu. one forward pass per prompt, around 0.2 ms.
"""

from __future__ import annotations

import sys
import time
from pathlib import Path

import torch

from modeling import encode_batch, load_router
from policy import Policy, Prototypes, decide

HERE = Path(__file__).parent
PROTOTYPES = HERE / "prototypes.json"


def route(prompts: list[str], policy_path: Path | None = None,
          prototypes_path: Path | None = None, device: str = "cpu"):
    model, tokenizer = load_router(HERE, device)
    policy = Policy.load(policy_path or HERE / "policy.json")
    prototypes = Prototypes.load(prototypes_path)
    tokens = encode_batch(tokenizer, prompts, model.config.max_len).to(device)
    return [decide(model, tokens[i:i + 1], prompts[i], policy, prototypes)
            for i in range(len(prompts))]


def add_category(name: str, examples: list[str], device: str = "cpu") -> None:
    """teach it a label it was never trained on. one forward pass, no optimizer."""
    model, tokenizer = load_router(HERE, device)
    tokens = encode_batch(tokenizer, examples, model.config.max_len).to(device)
    with torch.no_grad():
        projections = model(tokens)["proj"]
    prototypes = Prototypes.load(PROTOTYPES)
    prototypes.add(name, projections)
    prototypes.save(PROTOTYPES)
    print(f"added category {name!r} from {len(examples)} examples -> {PROTOTYPES}")


def show(decisions, prompts) -> None:
    for decision, prompt in zip(decisions, prompts, strict=True):
        flags = [k for k, v in decision.signals.items() if v]
        head = prompt if len(prompt) <= 58 else prompt[:55] + "..."
        print(f"{decision.tier:<9} {decision.domain:<20} c={decision.complexity} "
              f"d={decision.difficulty:.2f} {head}")
        print(f"          {decision.reason}"
              + (f" | signals={flags}" if flags else "")
              + (f" | category={decision.category}" if decision.category else ""))


DEMO = [
    "What's the capital of France?",
    "Write a short poem about a Python snake in the garden.",
    "Write a Python function that reverses a linked list in place.",
    "A train leaves at 3pm going 60mph, another at 4pm going 80mph, when do they meet?",
    "Summarize this email thread in two sentences.",
    "What are the latest news headlines today?",
    "Explain the difference between a mutex and a semaphore, with an example.",
]


def main() -> int:
    args = sys.argv[1:]
    if not args or args[0] in ("-h", "--help"):
        print(__doc__)
        return 0

    if args[0] == "--demo":
        started = time.perf_counter()
        decisions = route(DEMO)
        elapsed = time.perf_counter() - started
        show(decisions, DEMO)
        print(f"\n{len(DEMO)} prompts in {elapsed * 1000:.0f} ms "
              f"({elapsed * 1000 / len(DEMO):.2f} ms each, cpu, includes model load)")
        return 0

    if args[0] == "--add-category":
        if len(args) < 3:
            print("usage: python infer.py --add-category <name> <examples.txt>")
            return 1
        examples = [line.strip() for line in Path(args[2]).read_text().splitlines()
                    if line.strip()]
        add_category(args[1], examples)
        return 0

    policy_path = None
    if args[0] == "--policy":
        policy_path, args = Path(args[1]), args[2:]

    if args and args[0] == "--file":
        prompts = [line.strip() for line in Path(args[1]).read_text().splitlines()
                   if line.strip()]
    else:
        prompts = [" ".join(args)]

    decisions = route(prompts, policy_path,
                      PROTOTYPES if PROTOTYPES.is_file() else None)
    show(decisions, prompts)
    if len(prompts) == 1:
        print()
        print(decisions[0].as_line())
    return 0


if __name__ == "__main__":
    raise SystemExit(main())