import os

import gradio as gr

print("Starting Ada Space app...", flush=True)

BASE_MODEL_ID = os.environ.get("BASE_MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")
ADAPTER_ID = os.environ.get(
    "ADAPTER_ID",
    "IFthisisrealitynbds/ada-qwen-lora-adapter",
)
HF_TOKEN = os.environ.get("HF_TOKEN")

MODEL = None
TOKENIZER = None
TORCH = None


def load_model():
    global MODEL, TOKENIZER, TORCH

    if MODEL is not None and TOKENIZER is not None:
        return MODEL, TOKENIZER

    print("Loading Ada model...", flush=True)
    import torch
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer

    TORCH = torch
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32
    print(f"Using device: {device}", flush=True)

    TOKENIZER = AutoTokenizer.from_pretrained(
        BASE_MODEL_ID,
        token=HF_TOKEN,
        trust_remote_code=True,
    )
    if TOKENIZER.pad_token is None:
        TOKENIZER.pad_token = TOKENIZER.eos_token

    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        token=HF_TOKEN,
        torch_dtype=dtype,
        low_cpu_mem_usage=True,
        trust_remote_code=True,
    )
    MODEL = PeftModel.from_pretrained(base_model, ADAPTER_ID, token=HF_TOKEN)
    MODEL = MODEL.to(device)
    MODEL.eval()
    print("Ada model loaded.", flush=True)

    return MODEL, TOKENIZER


def build_messages(message, history):
    system_prompt = (
        "You are Ada, a helpful legal assistant. Be clear, practical, and careful. "
        "Do not pretend to be a lawyer. Encourage users to get professional legal advice "
        "for high-risk decisions."
    )

    messages = [{"role": "system", "content": system_prompt}]
    for item in history:
        if isinstance(item, dict):
            messages.append(item)
        else:
            user_message, assistant_message = item
            messages.append({"role": "user", "content": user_message})
            messages.append({"role": "assistant", "content": assistant_message})
    messages.append({"role": "user", "content": message})
    return messages


def respond(message, history):
    model, tokenizer = load_model()

    prompt = tokenizer.apply_chat_template(
        build_messages(message, history),
        tokenize=False,
        add_generation_prompt=True,
    )

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

    with TORCH.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=350,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )

    response = tokenizer.decode(
        output[0][inputs["input_ids"].shape[-1] :],
        skip_special_tokens=True,
    )
    return response.strip()


with gr.Blocks(title="Ada Legal Chatbot") as demo:
    gr.Markdown("# Ada Legal Chatbot")
    gr.ChatInterface(
        fn=respond,
        type="messages",
        cache_examples=False,
        run_examples_on_click=False,
        examples=[
            "What should I check before signing a tenancy agreement?",
            "Can my landlord evict me without notice?",
            "What can I do if repairs are not being done?",
        ],
    )


if __name__ == "__main__":
    print("Launching Gradio...", flush=True)
    demo.queue(default_concurrency_limit=1).launch(ssr_mode=False, show_error=True)