import streamlit as st
from sentence_transformers import SentenceTransformer, util
from groq import Groq
import os

# 1. Page Configuration & Professional High-Contrast Light Theme
st.set_page_config(
    page_title="ZeroAi Assistant",
    page_icon="⚡",
    layout="centered"
)

st.markdown("""
    <style>
    /* Premium Clean Light Background */
    .stApp {
        background: radial-gradient(circle at 10% 20%, rgba(16, 185, 129, 0.04) 0%, transparent 40%),
                    radial-gradient(circle at 90% 80%, rgba(59, 130, 246, 0.04) 0%, transparent 40%),
                    #f8fafc;
        color: #0f172a;
    }
    
    /* Elegant Clean Header */
    .chat-header {
        text-align: center;
        padding: 20px;
        background: #ffffff;
        border: 1px solid #e2e8f0;
        border-radius: 16px;
        margin-bottom: 30px;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.05);
    }
    .chat-header h1 {
        background: linear-gradient(to right, #0f172a, #10b981, #2563eb);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-weight: 800;
        font-size: 2.4rem;
        margin: 0;
    }
    
    /* High-Contrast Crisp Chat Bubbles */
    .stChatMessage {
        background-color: #ffffff !important;
        border: 1px solid #e2e8f0 !important;
        border-radius: 14px !important;
        margin-bottom: 14px !important;
        box-shadow: 0 1px 3px 0 rgba(0, 0, 0, 0.05);
        color: #0f172a !important;
    }
    
    /* Left/Right Border Identifiers for visibility */
    div[data-chat-message-user="true"] { 
        border-right: 4px solid #10b981 !important; 
        background-color: #f0fdf4 !important;
    }
    div[data-chat-message-assistant="true"] { 
        border-left: 4px solid #2563eb !important; 
    }
    
    /* Smooth, high-contrast text styling for standard paragraphs */
    .stChatMessage p, .stChatMessage div {
        color: #1e293b !important;
        font-size: 0.95rem !important;
        line-height: 1.6 !important;
    }

    /* Fixed Bottom Input Area Customization */
    div[data-testid="stChatInput"] {
        background-color: #ffffff !important;
        border: 1px solid #cbd5e1 !important;
        border-radius: 16px !important;
        box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.05) !important;
    }
    div[data-testid="stChatInput"] textarea { 
        color: #0f172a !important; 
    }
    
    /* Custom style for sidebar layout */
    section[data-testid="stSidebar"] {
        background-color: #ffffff !important;
        border-right: 1px solid #e2e8f0 !important;
    }
    </style>
""", unsafe_allow_html=True)

# 2. Predefined Responses for Basic Chatbot Logic
PREDEFINED_RESPONSES = {
    "hello": "Hello! I am ZeroAi, your everyday assistant. How can I help you today? 👋",
    "hi": "Hi there! I'm ZeroAi. Ready to answer regular questions or map out machine learning tasks.",
    "who are you": "I am ZeroAi, a versatile AI assistant. I can handle basic daily tasks or recommend specialized machine learning models depending on what you need.",
    "help": "You can ask me simple questions, request text corrections, or describe a software problem to get open-source model recommendations!",
    "clear": "To wipe our active chat session history, click the 'Reset Workspace' button inside the sidebar panel."
}

# 3. Model Architecture Dataset
@st.cache_data
def get_model_universe():
    return [
        {"name": "distilbert-base-uncased-finetuned-sst-2-english", "category": "Sentiment Analysis / Text Classification", "speed": "95/100", "accuracy": "91%", "size": "268 MB", "desc": "Fastest choice for production text sentiment analysis."},
        {"name": "cardiffnlp/twitter-roberta-base-sentiment-latest", "category": "Sentiment Analysis / Text Classification", "speed": "72/100", "accuracy": "95%", "size": "499 MB", "desc": "Highly accurate on slang, emojis, and social media layout nuances."},
        {"name": "prajjwal1/bert-tiny", "category": "Sentiment Analysis / Text Classification", "speed": "99/100", "accuracy": "76%", "size": "17.8 MB", "desc": "Ultra-lightweight footprint optimized for mobile or edge deployment."},
        {"name": "facebook/bart-large-cnn", "category": "Summarization", "speed": "48/100", "accuracy": "96%", "size": "1.63 GB", "desc": "Gold standard for generating coherent, abstractive long summaries."},
        {"name": "sshleifer/distilbart-cnn-12-6", "category": "Summarization", "speed": "82/100", "accuracy": "90%", "size": "1.20 GB", "desc": "Great balance between low latency and content recall."},
        {"name": "dbmdz/bert-large-cased-finetuned-conll03-english", "category": "Named Entity Recognition (NER)", "speed": "60/100", "accuracy": "97%", "size": "1.33 GB", "desc": "Flawless detection of organizations, people, locations, and data keys."},
        {"name": "elastic/distilbert-base-cased-finetuned-conll03-english", "category": "Named Entity Recognition (NER)", "speed": "94/100", "accuracy": "89%", "size": "261 MB", "desc": "Lean setup for high-volume real-time token text streams parsing."},
        {"name": "Helsinki-NLP/opus-mt-en-de", "category": "Translation", "speed": "85/100", "accuracy": "92%", "size": "298 MB", "desc": "Highly reliable local translation model for European language shifts."},
        {"name": "facebook/m2m100_418M", "category": "Translation", "speed": "55/100", "accuracy": "90%", "size": "1.84 GB", "desc": "Can translate directly between 100 languages without routing through English."},
        {"name": "deepset/roberta-base-squad2", "category": "Question Answering", "speed": "74/100", "accuracy": "93%", "size": "496 MB", "desc": "Excellent for context search engines extracting snippets from user documentation databases."},
        {"name": "Intel/dynamic_tinybert_squad2", "category": "Question Answering", "speed": "92/100", "accuracy": "84%", "size": "114 MB", "desc": "Accelerated quantization format ensuring nimble answers over shared networks."},
        {"name": "Qwen/Qwen2.5-Coder-7B-Instruct", "category": "Code Generation & Syntax Design", "speed": "68/100", "accuracy": "94%", "size": "14.0 GB", "desc": "State-of-the-art weights handling programming scripts and algorithm builds."},
        {"name": "HuggingFaceTB/SmolLM2-1.3B-Instruct", "category": "General Text Generation & Instructions", "speed": "96/100", "accuracy": "82%", "size": "2.6 GB", "desc": "Compact local chat companion running beautifully on minimal consumer hardware setups."}
    ]

@st.cache_resource
def init_local_embedder():
    return SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

model_universe = get_model_universe()
embedder = init_local_embedder()

categories_list = list(set([m["category"] for m in model_universe]))
category_vectors = embedder.encode(categories_list, convert_to_tensor=True)

# Crash-proof API environment lookup
GROQ_KEY = os.environ.get("GROQ_API_KEY") or "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"
groq_client = Groq(api_key=GROQ_KEY)

# 4. Interface Header Display
st.markdown("""
    <div class="chat-header">
        <h1>ZeroAi Portal</h1>
        <p style="color: #64748b; font-size: 0.95rem; margin-top: 5px; font-weight: 500;">
            ⚡ Smart Assistant & Model Recommendation Hub
        </p>
    </div>
""", unsafe_allow_html=True)

# Sidebar layout elements
with st.sidebar:
    st.markdown("### Controls")
    if st.button("Reset Workspace"):
        st.session_state.chat_history = [
            {"role": "assistant", "content": "Hello! I am ZeroAi. Ask me any basic question, or describe an engineering problem to explore tailored model suggestions. 🚀"}
        ]
        st.rerun()

# 5. Chat History Initialization
if "chat_history" not in st.session_state:
    st.session_state.chat_history = [
        {"role": "assistant", "content": "Hello! I am ZeroAi. Ask me any basic question, or describe an engineering problem to explore tailored model suggestions. 🚀"}
    ]

# Render persistent historical elements
for chat in st.session_state.chat_history:
    avatar_char = "⚡" if chat["role"] == "assistant" else "👤"
    with st.chat_message(chat["role"], avatar=avatar_char):
        st.write(chat["content"])

# 6. Stream Live Chat Interactions
if prompt := st.chat_input("Message ZeroAi..."):
    with st.chat_message("user", avatar="👤"):
        st.write(prompt)
    st.session_state.chat_history.append({"role": "user", "content": prompt})
    
    with st.chat_message("assistant", avatar="⚡"):
        placeholder = st.empty()
        clean_query = prompt.lower().strip().replace("?", "")
        
        # Branch 1: Predefined Match Check
        if clean_query in PREDEFINED_RESPONSES:
            predefined_ans = PREDEFINED_RESPONSES[clean_query]
            placeholder.markdown(predefined_ans)
            st.session_state.chat_history.append({"role": "assistant", "content": predefined_ans})
            
        # Branch 2: Advanced AI Processing
        else:
            # Semantic search check to flag model architectural intents
            prompt_vector = embedder.encode(prompt, convert_to_tensor=True)
            search_match = util.semantic_search(prompt_vector, category_vectors, top_k=1)
            best_similarity_score = search_match[0][0]['score']
            identified_arena = categories_list[search_match[0][0]['corpus_id']]
            
            # Context builder injection
            rec_context_str = ""
            if best_similarity_score > 0.45:
                matched_models = [m for m in model_universe if m["category"] == identified_arena]
                rec_context_str = f"\n[INTERNAL KNOWLEDGE MATCHED MODEL SUITE]\nDOMAIN FIELD: {identified_arena}\nTOP CHOSEN SELECTIONS:\n"
                for m in matched_models[:2]:
                    rec_context_str += f"- Name: {m['name']} (Speed rank: {m['speed']}, Accuracy tier: {m['accuracy']}) -> Description: {m['desc']}\n"
            
            expert_system_prompt = f"""
            You are ZeroAi, a friendly, ultra-intelligent, and clear AI assistant chatbot.
            The user prompt is: "{prompt}"
            
            Retrieved context (if applicable):
            {rec_context_str}
            
            Instructions:
            1. If the context maps to specific machine learning models, show the user these recommendations cleanly, highlighting their speed, accuracy, and size, and add a brief 4-line python code snippet using `transformers` to initialize it.
            2. If it's a general topic query (like cooking, essay outlines, email drafts, basic math, or definitions), completely ignore the models context and write a beautifully clear, direct, high-contrast light-mode readable response to satisfy their request.
            """
            
            try:
                response_stream = groq_client.chat.completions.create(
                    model="llama-3.3-70b-versatile",
                    messages=[{"role": "system", "content": expert_system_prompt}],
                    temperature=0.4,
                    stream=True
                )
                
                complete_text = ""
                for chunk in response_stream:
                    if chunk.choices[0].delta.content is not None:
                        complete_text += chunk.choices[0].delta.content
                        placeholder.markdown(complete_text + " ▌")
                        
                placeholder.markdown(complete_text)
                st.session_state.chat_history.append({"role": "assistant", "content": complete_text})
                
            except Exception as e:
                placeholder.error(f"ZeroAi linkage error. Technical parameters logs: {str(e)}")