import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
import sys
import argparse

def main():
    parser = argparse.ArgumentParser(description="Quintus Interactive Chat")
    parser.add_argument("--model_path", type=str, default="iamrahulreddy/Quintus", help="Model repo ID or local weights directory")
    parser.add_argument("--trust_remote_code", action="store_true", help="Allow custom code from the model repository.")
    args = parser.parse_args()

    model_path = args.model_path
    print(f"Loading Quintus from {model_path}...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=args.trust_remote_code)
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            device_map="auto",
            dtype=torch.float16,
            trust_remote_code=args.trust_remote_code
        )
    except Exception as e:
        print(f"Error loading model: {e}")
        print(f"Ensure '{model_path}' exists and contains the model weights.")
        sys.exit(1)

    # Defining stopping criteria
    stop_tokens = ["<|endoftext|>", "<|im_end|>"]
    eos_token_ids = [tokenizer.eos_token_id] if tokenizer.eos_token_id is not None else []
    for token in stop_tokens:
        t_id = tokenizer.convert_tokens_to_ids(token)
        if t_id is not None and t_id not in eos_token_ids:
            eos_token_ids.append(t_id)

    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    conversation_history = [
        {"role": "system", "content": "You are Quintus, a highly capable AI assistant created by Muskula Rahul. You are helpful, precise, and logically sound."}
    ]

    print()
    print("Quintus Chat (type 'quit' to exit)")
    print()

    while True:
        try:
            user_input = input("You: ").strip()
            if user_input.lower() in ["quit", "exit"]:
                print("\nGoodbye!")
                break
            if not user_input:
                continue

            conversation_history.append({"role": "user", "content": user_input})

            prompt = tokenizer.apply_chat_template(
                conversation_history,
                tokenize=False,
                add_generation_prompt=True
            )

            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

            print("Quintus: ", end="", flush=True)

            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=512,
                    temperature=0.7,
                    top_p=0.9,
                    do_sample=True,
                    streamer=streamer,
                    pad_token_id=tokenizer.eos_token_id,
                    eos_token_id=eos_token_ids
                )

            # Extract response for history
            generated_ids = outputs[0][inputs.input_ids.shape[-1]:]
            assistant_response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
            conversation_history.append({"role": "assistant", "content": assistant_response})
            print()

        except KeyboardInterrupt:
            print("\n\nGoodbye!")
            break

if __name__ == "__main__":
    main()