import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
import torch
from huggingface_hub import login
import os

# Force authentication with your HF token (secret in Space settings)
login(token=os.getenv("HF_TOKEN"))

# ================= CACHE THE MODEL =================
@st.cache_resource
def load_model():
    model_id = "ammoncoder123/IPTchatbotModel1-1.7B"  # Your correct repo

    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.float16
    )

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        quantization_config=quantization_config,
        device_map="auto",
        torch_dtype=torch.float16,
        trust_remote_code=True
    )

    return pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        max_new_tokens=300,
        temperature=0.7,
        do_sample=True,
        top_p=0.9
    )

pipe = load_model()

# ==================== CHAT INTERFACE ====================
st.title("IPT Chatbot Assistance")
st.info("Answers may vary — please verify important facts.")

# Display chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# User input
if prompt := st.chat_input("Ask about Industrial Practical Training..."):
    # Add user message to history and display
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Generate assistant response
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            # Strong system prompt (customized for your IPT dataset)
            system_prompt = """
            You are a helpful assistant for engineering and ICT students in Tanzania who are preparing for or doing Industrial Practical Training (IPT), also known as Industrial Attachment.
            IPT means Industrial Practical Training — a mandatory work placement where students gain real-world experience in companies related to their field of study.
            Always answer questions about:
            - What IPT is
            - How to do IPT (logbook, daily/weekly reports, technical report, presentation)
            - Placement suggestions for different engineering fields (ICT, Mechatronics, Electrical, Mechanical, Civil, Biomedical, etc.)
            - Choosing IPT centers/companies
            - Tips for success in IPT
            - Any other directly related IPT topic
            If the question is clearly unrelated to IPT (e.g., politics, sports, personal life), politely reply:
            "Sorry, I can only help with questions about Industrial Practical Training (IPT). Please ask something related to IPT, logbook, placement, or reports."
            For placement suggestions (e.g., for Mechatronics, Electrical, ICT), give practical, realistic company types or industries in Tanzania that match the field.
            Be concise, accurate, and helpful.
            """

            # Build messages: system prompt first, then user prompt
            chat_messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt}
            ]

            # Generate response
            outputs = pipe(
                chat_messages,
                max_new_tokens=300,
                temperature=0.7,
                do_sample=True,
                top_p=0.9
            )

            # Extract and clean the generated text
            response = outputs[0]["generated_text"]
            if isinstance(response, str) and response.startswith(prompt):
                response = response[len(prompt):].strip()

            # Show the response
            st.markdown(response)

    # Save assistant response to history
    st.session_state.messages.append({"role": "assistant", "content": response})

# Optional: Clear conversation button
if st.button("Clear Conversation"):
    st.session_state.messages = []
    st.rerun()