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)
# Improved custom CSS
st.markdown("""
    <style>
    .main {
        background: linear-gradient(to right, #1f4037, #99f2c8);
        padding: 2rem;
        font-family: 'Segoe UI', sans-serif;
    }
    .stButton>button {
        background-color: #ffffff10;
        color: white;
        font-weight: 600;
        border-radius: 10px;
        padding: 0.6rem 1rem;
        transition: all 0.3s ease;
        border: 1px solid white;
    }
    .stButton>button:hover {
        background-color: #ffffff30;
        color: white;
        border-color: #fff;
    }
    h1, h3, p, label {
        color: white;
        text-align: center;
    }
    .chat-bubble-user {
        background-color: #ffffff25;
        padding: 0.75rem 1rem;
        border-radius: 1rem;
        margin-bottom: 0.5rem;
        color: #fff;
        font-weight: 500;
        text-align: left;
        width: fit-content;
        max-width: 90%;
        align-self: flex-end;
    }
    .chat-bubble-bot {
        background-color: #ffffff15;
        padding: 0.75rem 1rem;
        border-radius: 1rem;
        margin-bottom: 0.5rem;
        color: #fff;
        text-align: left;
        width: fit-content;
        max-width: 90%;
        align-self: flex-start;
    }
    .chat-container {
        display: flex;
        flex-direction: column;
        gap: 0.5rem;
        margin-top: 2rem;
    }
    </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 seasoned Machine Learning mentor with {experience_label} years of hands-on expertise. Your teaching style is friendly, clear, and approachable. Follow these strict guidelines:

        1. Only respond to questions directly related to machine learning programming — including its libraries, tools, and frameworks.
        2. If asked about anything outside machine learning, reply with: "I specialize only in Machine learning programming. This appears to be a non-machine learning topic."
        3. Do not offer help on topics unrelated to machine learning.
        4. Keep your explanations beginner-friendly when needed, focusing on clarity and real-world application.
        5. Use practical code examples and scenarios to reinforce learning.
        6. For complex or advanced topics, assume the learner has foundational knowledge of machine learning concepts."""
        
       
        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)

if st.session_state[PAGE_KEY]:
    st.markdown('<div class="chat-container">', unsafe_allow_html=True)
    for user, bot in st.session_state[PAGE_KEY]:
        st.markdown(f'<div class="chat-bubble-user">👤 <strong>You:</strong> {user}</div>', unsafe_allow_html=True)
        st.markdown(f'<div class="chat-bubble-bot">🧑‍🏫 <strong>Mentor:</strong> {bot}</div>', unsafe_allow_html=True)
    st.markdown('</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")