#!/usr/bin/env python
"""PALIMPSESTE — Interactive chat with multi-turn memory and live learning.

Features:
  - Multi-turn conversation: the model remembers previous Q/A turns
  - Live learning: type "apprends que <question> = <reponse>" to teach instantly
  - Fuzzy matching: typos and near-misses are matched to known questions
  - Streaming: tokens appear as they are generated
  - Graceful fallback: unknown questions get a helpful "teach me" prompt

Usage:
  python examples/chat.py --model ./chat_model
  python examples/chat.py --model ./chat_model --temperature 0.3 --stream
  python examples/chat.py --model user/palimpseste-chat  # from HF Hub

Commands:
  apprends que <question> = <reponse>   — teach a new Q/A pair (instant, O(1))
  teach <question> = <answer>           — same, in English
  /transcript                           — show the full conversation
  /reset                                — clear conversation history
  /stats                                — show model statistics
  quit / exit / bye / au revoir         — exit
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from palimseste.hf import HFPalimpsesteLM
from palimseste.chat import Conversation, FALLBACK_RESPONSE


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Chat with a trained PALIMPSESTE model (multi-turn, live learning).",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("--model", "-m", type=str, required=True,
                   help="path to a saved model dir, or an HF Hub repo id")
    p.add_argument("--temperature", "-t", type=float, default=0.3,
                   help="sampling temperature (0 = greedy, higher = more random)")
    p.add_argument("--max-tokens", type=int, default=200,
                   help="max tokens per response")
    p.add_argument("--seed", type=int, default=None,
                   help="random seed for reproducible responses")
    p.add_argument("--stream", action="store_true", default=True,
                   help="stream tokens as they are generated (default: on)")
    p.add_argument("--no-stream", dest="stream", action="store_false",
                   help="disable streaming")
    p.add_argument("--history", type=int, default=3,
                   help="number of past turns to remember (default: 3)")
    p.add_argument("--fuzzy", type=float, default=0.75,
                   help="fuzzy match threshold 0..1 (default: 0.75)")
    p.add_argument("--qa-file", type=str, default=None,
                   help="JSON file with [{q,a}] to pre-register for fuzzy matching")
    return p.parse_args()


def load_model(model_path: str) -> HFPalimpsesteLM:
    p = Path(model_path)
    if p.exists() and p.is_dir():
        return HFPalimpsesteLM.from_pretrained(p)
    try:
        from huggingface_hub import snapshot_download
        local = snapshot_download(repo_id=model_path, repo_type="model")
        return HFPalimpsesteLM.from_pretrained(local)
    except Exception as e:
        sys.exit(f"could not load model from '{model_path}': {e}")


def load_known_pairs(args) -> list[tuple[str, str]]:
    """Load Q/A pairs for fuzzy matching registration."""
    if args.qa_file:
        import json
        with open(args.qa_file, "r", encoding="utf-8") as f:
            data = json.load(f)
        return [(d["q"], d["a"]) for d in data]
    # try the built-in corpus
    try:
        from corpus_chat import get_corpus
        return get_corpus()
    except ImportError:
        return []


BANNER = r"""
 ____  _               _ _____            _
|  _ \| | __ _ _   _  | |  ___| ___   __| | __ _
| |_) | |/ _` | | | | | | |_   / _ \ / _` |/ _` |
|  __/| | (_| | |_| | | |  _| | (_) | (_| | (_| |
|_|   |_|\__,_|\__, | |_|_|    \___/ \__,_|\__,_|
               |___/
           Hypervectorial Cortex · Multi-turn Chat
"""


def main() -> None:
    args = parse_args()
    print(f"loading model from {args.model} ...", file=sys.stderr)
    lm = load_model(args.model)
    print(f"loaded: D={lm.config.D:,}  |M|={len(lm.mem):,}  "
          f"vocab={lm.config.vocab_size}", file=sys.stderr)

    # build conversation
    conv = Conversation(
        model=lm,
        max_history_turns=args.history,
        fuzzy_threshold=args.fuzzy,
        learn_live=True,
    )
    # register known questions for fuzzy matching
    known = load_known_pairs(args)
    if known:
        conv.register_questions(known)
        print(f"registered {len(known)} Q/A pairs for fuzzy matching",
              file=sys.stderr)

    print(BANNER)
    print("Multi-turn chat with live learning. Type a message and press Enter.")
    print()
    print("Commands:")
    print("  apprends que <question> = <reponse>  — teach instantly (O(1))")
    print("  /transcript  /reset  /stats          — conversation controls")
    print("  quit                                  — exit")
    print(f"(temperature={args.temperature}, history={args.history} turns, "
          f"fuzzy={args.fuzzy}, stream={args.stream})")
    print("-" * 60)

    while True:
        try:
            user = input("\nYou> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n\nPalimpseste> Au revoir!")
            break
        if not user:
            continue

        # slash commands
        if user.startswith("/"):
            cmd = user.lower().strip("/")
            if cmd in ("quit", "exit", "bye"):
                print("Palimpseste> Au revoir!")
                break
            elif cmd == "transcript":
                print("\n--- Conversation Transcript ---")
                print(conv.get_transcript())
                print("--- End ---")
                continue
            elif cmd == "reset":
                conv.reset()
                print("Palimpseste> (conversation history cleared)")
                continue
            elif cmd == "stats":
                s = lm.stats()
                print(f"  D={s['D']:,}  |M|={s['n_traces']:,}  "
                      f"vocab={s['vocab_size']}  turns={conv.turn_count}")
                continue
            else:
                print(f"unknown command: /{cmd}")
                continue

        if user.lower() in ("quit", "exit", "bye", "au revoir"):
            print("Palimpseste> Au revoir! A bientot.")
            break

        # check for teach command
        if conv.try_parse_teach(user):
            # the teach was already processed and recorded in history
            # print the confirmation (last palimpseste turn)
            last = conv.history[-1]
            print(f"Palimpseste> {last.text}")
            continue

        # normal response
        if args.stream:
            print("Palimpseste> ", end="", flush=True)
            for token_text, _ in conv.respond_streaming(
                user, max_new_tokens=args.max_tokens,
                temperature=args.temperature, seed=args.seed
            ):
                print(token_text, end="", flush=True)
            print()  # newline after streaming
        else:
            resp = conv.respond(user, max_new_tokens=args.max_tokens,
                                temperature=args.temperature, seed=args.seed)
            print(f"Palimpseste> {resp}")


if __name__ == "__main__":
    main()