import argparse
import os
os.environ["PYTHONUTF8"] = "1"  # Keep Windows encoding clean

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

# 1. Hardware setup
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Running inference on: {DEVICE}")

# 2. Define our local path and the base model we used
BASE_MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
ADAPTER_DIR = "./fine_tuned_gen_model"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)

print("Loading base model weights...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
).to(DEVICE)

print("Merging fine-tuned Juspay adapters...")
# This layers your custom weights on top of the base model
model = PeftModel.from_pretrained(base_model, ADAPTER_DIR).to(DEVICE)
model.eval()  # Put model in evaluation mode


def generate_answer(user_query: str) -> str:
    prompt = f"User: {user_query}\nAssistant:"
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=150,
            temperature=0.3,
            do_sample=True,
            top_p=0.95,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return generated_text.replace(prompt, "").strip()


def run_cli() -> None:
    print("\n🚀 Juspay Interview Bot Initialized! (Type 'quit' to exit)")
    print("-" * 50)
    while True:
        user_query = input("\nCandidate Question: ")
        if user_query.strip().lower() == "quit":
            break
        answer = generate_answer(user_query)
        print(f"\nAI Response:\n{answer}")
        print("-" * 50)


def run_gradio() -> None:
    import gradio as gr

    with gr.Blocks() as demo:
        gr.Markdown("# Juspay Interview Bot")
        gr.Markdown(
            "Ask the fine-tuned Qwen model your interview questions and get an instant response."
        )
        question = gr.Textbox(
            label="Candidate Question",
            placeholder="Enter your interview question here...",
            lines=4,
        )
        answer = gr.Textbox(label="AI Response", lines=8)
        submit = gr.Button("Generate")
        submit.click(fn=generate_answer, inputs=question, outputs=answer)

    demo.launch(share=False)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run the Juspay Interview Bot")
    parser.add_argument(
        "--interface",
        choices=["cli", "gradio"],
        default="gradio",
        help="Choose interface mode",
    )
    args = parser.parse_args()

    if args.interface == "cli":
        run_cli()
    else:
        run_gradio()