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

# ================= AUTHENTICATION =================
# Ensure this is flush to the left margin
hf_token = os.getenv("HF_TOKEN")
if hf_token:
    try:
        login(token=hf_token)
    except Exception:
        # Silently fail if network is temporarily down
        pass

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

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

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    if tokenizer.chat_template is None:
        tokenizer.chat_template = "{% for message in messages %}{{'<|' + message['role'] + '|>' + '\n' + message['content'] + '</s>\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|assistant|>\n' }}{% endif %}"

    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)

pipe = load_model()

# CHAT INTERFACE 
st.title("Industrial Practical Training Chatbot")
st.info("Ask about logbooks, placement, or report writing.")

if "messages" not in st.session_state:
    st.session_state.messages = []

# Display history
for message in st.session_state.messages:
    if message["role"] != "system":
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

# User input
if prompt := st.chat_input("Ask about IPT..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        with st.spinner("Analyzing IPT guidelines..."):
            system_message = {
                "role": "system",
                "content": """You are the "IPT Master Assistant," a specialized AI coordinator for engineering and ICT students in Tanzania undergoing Industrial Practical Training (IPT). Your goal is to provide accurate, encouraging, and practical advice based on official training guidelines.

### CORE KNOWLEDGE & RULES:
1. DAILY LOGBOOK:
   - Entries must include: Major tasks performed, time spent (in hours), and any absences with reasons.
   - Style: Keep it simple and clear (3-5 sentences). 
   - Importance: It accounts for 10% of the final grade.

2. WEEKLY REPORTS:
   - Focus on ONE major technical task from the week (Maintenance, Repair, or Production).
   - Use the "Detailed Description" page. Be concise but technical.

3. PLACEMENT STRATEGY:
   - Placement Must Match Department: Always advise students to choose a center that correlates with their field.
   - ICT Recommendations: Suggest IT departments in firms like TTCL, Halotel, Vodacom, or e-GA.

4. BEHAVIORAL GUIDELINES:
   - Polite & Professional: If unrelated to IPT, politely decline.
   - Fact-Checking: 100% honesty is required.

### TONE:
Academic mentor. Refer to Tanzanian institutions (ATC, DIT, MUST, UDSM)."""
            }

            # Standardized indentation for input_messages
            input_messages = [system_message] + st.session_state.messages[-3:]

            formatted_prompt = pipe.tokenizer.apply_chat_template(
                input_messages, 
                tokenize=False, 
                add_generation_prompt=True
            )

            outputs = pipe(
                formatted_prompt,
                max_new_tokens=300,
                temperature=0.7,
                do_sample=True,
                top_p=0.9,
                return_full_text=False 
            )

            response = outputs[0]["generated_text"].strip()
            st.markdown(response)

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

if st.button("Clear History"):
    st.session_state.messages = []
    st.rerun()