import streamlit as st
import os
from storage import ensure_dirs
from rag_core import query_documents, run_openai_agent, finalize_answer

# ---------------------------------------------------------------------------------
# STARTUP
# ---------------------------------------------------------------------------------
ensure_dirs()

# ---------------------------------------------------------------------------------
# PAGE CONFIGURATION
# ---------------------------------------------------------------------------------
st.set_page_config(
    page_title="DevRAG Assistant",
    page_icon="💻",
    layout="centered",
    initial_sidebar_state="expanded",
)

# ---------------------------------------------------------------------------------
# CUSTOM CSS
# ---------------------------------------------------------------------------------
st.markdown(
    """
<style>
    .block-container {
        padding-top: 2rem !important;
        padding-bottom: 5rem !important;
        max-width: 850px !important;
    }
    h1 {
        font-weight: 700;
        font-size: 2.2rem;
        margin-bottom: 0px !important;
    }
    p.subtitle {
        opacity: 0.7;
        font-size: 1.05rem;
        margin-bottom: 1.5rem;
    }
    .stack-badge {
        display: inline-block;
        padding: 4px 12px;
        border-radius: 16px;
        font-size: 0.85rem;
        font-weight: 600;
        margin-right: 8px;
        background-color: rgba(255, 255, 255, 0.1);
        border: 1px solid rgba(255, 255, 255, 0.2);
    }
    .example-q {
        cursor: pointer;
        padding: 8px 12px;
        border-radius: 8px;
        background: rgba(255, 255, 255, 0.05);
        margin-bottom: 8px;
        font-size: 0.9rem;
        transition: 0.2s;
    }
    .example-q:hover {
        background: rgba(255, 255, 255, 0.1);
    }
</style>
""",
    unsafe_allow_html=True,
)

# ---------------------------------------------------------------------------------
# STATE MANAGEMENT
# ---------------------------------------------------------------------------------
if "messages" not in st.session_state:
    st.session_state["messages"] = [
        {
            "role": "assistant",
            "content": "Hello! I am your focused programming assistant for **JavaScript, Python, and PHP**. How can I help you code today?",
            "sources": [],
        }
    ]

# ---------------------------------------------------------------------------------
# SIDEBAR
# ---------------------------------------------------------------------------------
with st.sidebar:
    st.markdown("### ⚙️ Settings")
    api_key_env = os.getenv("OPENAI_API_KEY", "")
    api_key_input = st.text_input(
        "OpenAI API Key",
        value=api_key_env,
        type="password",
        placeholder="sk-...",
    )

    if api_key_input:
        os.environ["OPENAI_API_KEY"] = api_key_input
        st.session_state["api_key"] = api_key_input
    else:
        st.session_state["api_key"] = None

    st.divider()
    
    st.markdown("### 📚 Supported Stacks")
    st.markdown("- **Python**: core, libraries, logic")
    st.markdown("- **JavaScript**: node, web, frameworks")
    st.markdown("- **PHP**: general, laravel patterns")
    
    st.divider()
    
    st.markdown("### 💡 Example Questions")
    examples = [
        "How do Python list comprehensions work?",
        "Explain JavaScript async/await with an example.",
        "What is a Service Class pattern in PHP?",
        "Difference between Python lists and JS arrays."
    ]
    for q in examples:
        if st.button(q, key=f"ex_{q}", use_container_width=True):
            st.session_state["clicked_example"] = q
            st.rerun()

# ---------------------------------------------------------------------------------
# MAIN INTERFACE
# ---------------------------------------------------------------------------------
st.markdown(
    """
<div style="text-align: center; margin-bottom: 2rem;">
    <h1>💻 DevRAG Assistant</h1>
    <p class='subtitle'>Focused knowledge RAG for modern stacks</p>
    <div>
        <span class="stack-badge">🐍 Python</span>
        <span class="stack-badge">🟨 JavaScript</span>
        <span class="stack-badge">🐘 PHP</span>
    </div>
</div>
""",
    unsafe_allow_html=True,
)

# Display chat messages
for msg in st.session_state["messages"]:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])
        if msg.get("sources"):
            with st.expander("📚 Referenced Repository Samples", expanded=False):
                for s in msg["sources"]:
                    st.markdown(f"- `{s['citation']}`")

# Handle input
prompt = st.chat_input("Ask a question about JS, Python, or PHP...")

# Handle clicked example from sidebar
if st.session_state.get("clicked_example"):
    prompt = st.session_state.pop("clicked_example")

if prompt:
    if not st.session_state.get("api_key"):
        st.error("Please enter your OpenAI API Key in the sidebar.")
    else:
        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 codebases & thinking..."):
                try:
                    retrieved = query_documents(prompt, top_k=5)
                    raw_response, tool_used = run_openai_agent(
                        query=prompt,
                        context_items=retrieved,
                        api_key=st.session_state["api_key"],
                    )
                    final_text, sources = finalize_answer(
                        raw_response, retrieved, tool_used=tool_used
                    )

                    st.markdown(final_text)
                    if sources:
                        with st.expander("📚 Referenced Repository Samples", expanded=False):
                            for s in sources:
                                st.markdown(f"- `{s['citation']}`")

                    st.session_state.messages.append(
                        {
                            "role": "assistant",
                            "content": final_text,
                            "sources": sources,
                        }
                    )
                except Exception as e:
                    st.error(f"Error: {e}")