import streamlit as st
import os
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langchain_core.messages import HumanMessage, SystemMessage

# Set Hugging Face tokens
hf = os.getenv('HF_TOKEN')
os.environ['HUGGINGFACEHUB_API_TOKEN'] = hf
os.environ['HF_TOKEN'] = hf

# --- Page Configuration ---
st.set_page_config(page_title="Machine Learning Mentor", layout="wide")

# --- Custom CSS ---
st.markdown("""
    <style>
    body {
        font-family: 'Segoe UI', sans-serif;
        background: linear-gradient(135deg, #2c003e, #0f9b8e);
    }
    .main {
        padding: 2rem;
    }
    .stTextInput>div>div>input {
        color: #ffffff !important;
        background-color: #00000010;
        border: 1px solid #ffffff30;
    }
    .chat-bubble {
        padding: 1rem;
        margin: 0.5rem 0;
        border-radius: 12px;
        max-width: 80%;
    }
    .user-bubble {
        background-color: #4b0082;
        color: white;
        margin-left: auto;
        text-align: right;
    }
    .mentor-bubble {
        background-color: #00bfa6;
        color: white;
        margin-right: auto;
    }
    .stButton>button {
        background: #ffffff20;
        border: 1px solid #ffffff50;
        color: white;
        font-weight: bold;
        border-radius: 10px;
        transition: 0.3s ease;
    }
    .stButton>button:hover {
        background: #ffffff40;
    }
    </style>
""", unsafe_allow_html=True)

# --- Title ---
st.title("🤖 Machine Learning Mentor")

# --- Sidebar Preferences ---
st.sidebar.title("Mentor Preferences")
experience_label = st.sidebar.selectbox("Select your experience level:", ["Beginner", "Intermediate", "Experienced"])

# --- Initialize Model ---
ml_model_skeleton = HuggingFaceEndpoint(
    repo_id='Qwen/Qwen3-14B',
    provider='nebius',
    temperature=0.7,
    max_new_tokens=50,
    task='conversational'
)

ml_mentor = ChatHuggingFace(
    llm=ml_model_skeleton,
    repo_id='Qwen/Qwen3-14B',
    provider='nebius',
    temperature=0.7,
    max_new_tokens=50,
    task='conversational'
)

PAGE_KEY = "ml_chat_history"
if PAGE_KEY not in st.session_state:
    st.session_state[PAGE_KEY] = []

# --- Layout ---
col1, col2 = st.columns([3, 1])

# --- Chat Section ---
with col1:
    with st.form(key="chat_form"):
        user_input = st.text_input("Ask your question:")
        submit = st.form_submit_button("Send")

    if submit and user_input:
        system_prompt = (
            f"You are a machine learning mentor with {experience_label.lower()} experience. "
            f"Answer only machine learning questions in a friendly tone and within 150 words. "
            f"Politely inform if the question is out of scope."
        )
        messages = [SystemMessage(content=system_prompt), HumanMessage(content=user_input)]
        result = ml_mentor.invoke(messages)
        st.session_state[PAGE_KEY].append((user_input, result.content))

    st.subheader("🗨️ Chat History")
    for user, bot in st.session_state[PAGE_KEY]:
        st.markdown(f'<div class="chat-bubble user-bubble">{user}</div>', unsafe_allow_html=True)
        st.markdown(f'<div class="chat-bubble mentor-bubble">{bot}</div>', unsafe_allow_html=True)

# --- Mentor Tips Sidebar ---
with col2:
    st.markdown("### 💡 Tips from Mentor")
    st.info("Try asking about:\n- Regression vs Classification\n- Overfitting examples\n- Feature scaling\n- Model evaluation techniques")