"""
ViuAI-1M Chat Script.

Start an interactive chat session with a trained ViuAI model.

Usage:
    python scripts/chat.py --checkpoint output/final --tokenizer output/tokenizer
"""

import argparse
import os
import sys

# Add project root to path
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

from viuai.inference.chat import load_model_for_chat


def parse_args():
    parser = argparse.ArgumentParser(description="Chat with ViuAI-1M")
    parser.add_argument("--checkpoint", type=str, default="output/final",
                        help="Path to model checkpoint directory")
    parser.add_argument("--tokenizer", type=str, default="output/tokenizer",
                        help="Path to tokenizer directory")
    parser.add_argument("--device", type=str, default=None,
                        help="Device (cuda/cpu)")
    parser.add_argument("--temperature", type=float, default=0.8,
                        help="Sampling temperature")
    parser.add_argument("--max_tokens", type=int, default=128,
                        help="Maximum tokens to generate")
    return parser.parse_args()


def main():
    args = parse_args()

    checkpoint_path = os.path.join(project_root, args.checkpoint)
    tokenizer_path = os.path.join(project_root, args.tokenizer)

    print(f"\n  Loading ViuAI-1M...")
    print(f"  Checkpoint: {checkpoint_path}")
    print(f"  Tokenizer:  {tokenizer_path}")

    engine = load_model_for_chat(
        checkpoint_dir=checkpoint_path,
        tokenizer_dir=tokenizer_path,
        device=args.device,
    )

    # Start interactive chat
    engine.interactive()


if __name__ == "__main__":
    main()