#!/usr/bin/env python3
"""
Saturday AI — Interactive Chat & Text Generation CLI.

Run this script to interact with your from-scratch Saturday Language Model!

Usage:
    # 1. Quick interactive mode (trains a small model on built-in text in 5 seconds):
    python scripts/chat.py

    # 2. Interactive mode with custom prompt from command line:
    python scripts/chat.py --prompt "The future of artificial intelligence"

    # 3. Train on your own text file and chat:
    python scripts/chat.py --data my_data.txt --steps 300

    # 4. Load a saved checkpoint:
    python scripts/chat.py --checkpoint checkpoints/saturday_latest.json
"""

import sys
import os
import argparse
import time
import numpy as np

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

from saturday_numpy.config import SaturdayConfig
from saturday_numpy.tokenizer.character_tokenizer import CharacterTokenizer
from saturday_numpy.model.saturday import SaturdayModel
from saturday_numpy.training.loss import cross_entropy_loss
from saturday_numpy.training.optimizer import AdamW
from saturday_numpy.inference.generator import TextGenerator
from saturday_numpy.utils.checkpoint import load_checkpoint, save_checkpoint


DEFAULT_CORPUS = """
Saturday is a research-grade language model built from scratch to understand Transformer internals.
It uses a decoder-only architecture with Root Mean Square Normalization (RMSNorm), Rotary Position Embeddings (RoPE), SwiGLU Feed-Forward Networks, and Multi-Head Causal Self-Attention with Grouped Query Attention (GQA).
Language models predict the probability of the next token given all previous tokens in an autoregressive sequence.
Artificial intelligence models learn patterns, grammar, semantics, and reasoning from vast amounts of text data through gradient descent and backpropagation.
Deep learning models scale with parameters, datasets, and compute, enabling surprising capabilities in natural language understanding, code generation, reasoning, and conversation.
The future of AI is open, transparent, scalable, and built on fundamental mathematical principles.
"""


def train_quick_model(corpus_text: str, steps: int = 300, lr: float = 0.008):
    print("\n" + "=" * 60)
    print("  [1/2] Training Saturday-1M on text corpus...")
    print("=" * 60)
    
    tokenizer = CharacterTokenizer.build_from_text(corpus_text)
    seq_len = 128
    all_token_ids = np.array(tokenizer.encode(corpus_text), dtype=np.int32)
    
    # Create batched sequences of max 128 tokens
    batches = []
    for i in range(0, len(all_token_ids) - seq_len, seq_len // 2):
        batches.append(all_token_ids[i : i + seq_len + 1])
        if len(batches) >= 8:
            break
    if not batches:
        batches.append(all_token_ids)
        
    batch_arr = np.array(batches, dtype=np.int32)
    
    config = SaturdayConfig(
        vocab_size=tokenizer.vocab_size,
        hidden_size=128,
        intermediate_size=384,
        num_layers=4,
        num_attention_heads=4,
        num_key_value_heads=4,
        max_sequence_length=256,
        random_seed=42,
    )
    
    model = SaturdayModel(config)
    optimizer = AdamW(model=model, learning_rate=0.01, weight_decay=0.0)
    
    input_tokens = batch_arr[:, :-1]
    target_tokens = batch_arr[:, 1:]
    
    start_time = time.time()
    for step in range(1, steps + 1):
        logits = model.forward(input_tokens)
        loss, d_logits = cross_entropy_loss(logits, target_tokens)
        model.backward(d_logits)
        optimizer.step()
        
        if step % 50 == 0 or step == steps:
            elapsed = time.time() - start_time
            tok_sec = (step * input_tokens.size) / elapsed
            print(f"  Step {step:03d}/{steps} | Loss: {loss:.4f} | Speed: {tok_sec:.1f} tokens/sec")
            
    print(f"\n  [OK] Training complete in {time.time() - start_time:.2f}s! Final Loss: {loss:.4f}")
    
    # Save checkpoint
    os.makedirs("checkpoints", exist_ok=True)
    ckpt_path = "checkpoints/saturday_interactive.json"
    save_checkpoint(
        model=model,
        optimizer=optimizer,
        config=config,
        tokenizer=tokenizer,
        step=steps,
        train_tokens=steps * input_tokens.size,
        val_loss=float(loss),
        path=ckpt_path,
    )
    print(f"  [OK] Checkpoint saved to {ckpt_path}")
    
    return model, tokenizer, config


def interactive_chat_loop(model, tokenizer, config):
    generator = TextGenerator(model, tokenizer)
    
    print("\n" + "=" * 60)
    print("        SATURDAY AI — INTERACTIVE CHAT SESSION")
    print("=" * 60)
    print("  Model: Saturday (" + f"{model.count_parameters():,} parameters)")
    print("  Type a prompt and press Enter to generate response.")
    print("  Commands: 'exit' or 'quit' to end session.")
    print("=" * 60 + "\n")
    
    while True:
        try:
            prompt = input("\nUser > ").strip()
            if not prompt:
                continue
            if prompt.lower() in ("exit", "quit", "q"):
                print("\nEnding Saturday session. Goodbye!")
                break
                
            print("\nSaturday > ", end="", flush=True)
            
            # Check if prompt chars are in vocabulary
            try:
                tokenizer.encode(prompt)
            except KeyError as e:
                print(f"[Error: Prompt contains character not in vocabulary: {e}]")
                continue

            # Streamed/Autoregressive generation
            generated_text = generator.generate(
                prompt=prompt,
                max_new_tokens=100,
                temperature=0.2,  # Low temperature for clean English
                top_k=5,
                top_p=0.9,
                repetition_penalty=1.1,
            )
            
            # Print completion beyond prompt
            completion = generated_text[len(prompt):]
            for char in completion:
                print(char, end="", flush=True)
                time.sleep(0.015)  # typing effect
            print()

        except KeyboardInterrupt:
            print("\n\nSession interrupted. Goodbye!")
            break


def main():
    parser = argparse.ArgumentParser(description="Saturday AI Interactive Chat CLI")
    parser.add_argument("--data", type=str, help="Path to custom text file to train on")
    parser.add_argument("--checkpoint", type=str, help="Path to load saved checkpoint from")
    parser.add_argument("--prompt", type=str, help="One-shot prompt generation")
    parser.add_argument("--steps", type=int, default=200, help="Number of training steps")
    parser.add_argument("--lr", type=float, default=0.005, help="Learning rate")
    args = parser.parse_args()

    if args.checkpoint and os.path.exists(args.checkpoint):
        print(f"Loading checkpoint from {args.checkpoint}...")
        model, _, config, tokenizer, step = load_checkpoint(args.checkpoint)
        print(f"Loaded checkpoint at step {step}")
    else:
        corpus = DEFAULT_CORPUS
        if args.data and os.path.exists(args.data):
            with open(args.data, "r", encoding="utf-8") as f:
                corpus = f.read()
            print(f"Loaded custom text dataset from {args.data}")
            
        model, tokenizer, config = train_quick_model(corpus, steps=args.steps, lr=args.lr)

    if args.prompt:
        generator = TextGenerator(model, tokenizer)
        print(f"\nPrompt: {args.prompt}")
        output = generator.generate(prompt=args.prompt, max_new_tokens=100, temperature=0.2)
        print(f"Saturday: {output}")
    else:
        interactive_chat_loop(model, tokenizer, config)


if __name__ == "__main__":
    main()