import streamlit as st
import difflib
import requests
import datetime

# --- CONFIG ---
GROQ_API_KEY = st.secrets.get('GROQ_API_KEY', 'YOUR_GROQ_API_KEY')
BLACKBOX_API_KEY = st.secrets.get('BLACKBOX_API_KEY', 'YOUR_BLACKBOX_API_KEY')

PROGRAMMING_LANGUAGES = ["Python", "JavaScript", "TypeScript", "Java", "C++", "C#"]
SKILL_LEVELS = ["Beginner", "Intermediate", "Expert"]
USER_ROLES = ["Student", "Frontend Developer", "Backend Developer", "Data Scientist"]
EXPLANATION_LANGUAGES = ["English", "Spanish", "Chinese", "Urdu"]
EXAMPLE_QUESTIONS = [
    "What does this function do?",
    "How can I optimize this code?",
    "What are the potential bugs in this code?",
    "How does this algorithm work?",
    "What design patterns are used here?",
    "How can I make this code more readable?"
]

# --- API CALLS ---
def call_groq_api(prompt, model="llama3-70b-8192"):
    headers = {"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"}
    data = {"model": model, "messages": [{"role": "user", "content": prompt}]}
    response = requests.post("https://api.groq.com/openai/v1/chat/completions", json=data, headers=headers)
    if response.status_code == 200:
        return response.json()['choices'][0]['message']['content']
    else:
        return f"[Groq API Error] {response.text}"

def call_blackbox_agent(messages, model="gpt-4o"):
    """
    messages: list of dicts, e.g.
      [
        {"role": "system", "content": "You are a helpful coding assistant."},
        {"role": "user", "content": "Refactor this code: ..."}
      ]
    """
    url = "https://api.blackbox.ai/v1/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {BLACKBOX_API_KEY}"
    }
    data = {
        "model": model,
        "messages": messages
    }
    response = requests.post(url, headers=headers, json=data)
    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"]
    else:
        return call_groq_api(messages[-1]["content"])



# --- UTILS ---
def code_matches_language(code, language):
    if language.lower() in code.lower():
        return True
    return True 

def calculate_code_complexity(code):
    lines = code.count('\n') + 1
    return f"{lines} lines"

def get_inline_diff(original, modified):
    diff = difflib.unified_diff(
        original.splitlines(),
        modified.splitlines(),
        lineterm='',
        fromfile='Original',
        tofile='Refactored'
    )
    return '\n'.join(diff)

# --- STREAMLIT APP ---
st.set_page_config(page_title="Code Workflows", layout="wide")
st.title("CodeGenie")

# Navigation
page = st.sidebar.radio("Navigate", ["Home", "Code Workflow", "Semantic Search"])

if page == "Home":
    st.header("Welcome to the Code Genie!")
    st.markdown("""
    - **Full Code Workflow:** Complete code analysis pipeline with explanation, refactoring, review, and testing (powered by Groq/Blackbox)
    - **Semantic Search:** Ask natural language questions about your code and get intelligent answers
    """)
    st.info("Select a feature from the sidebar to get started.")

elif page == "Code Workflow":
    st.header("Full Code Workflow")
    code_input = st.text_area("Paste your code here", height=200)
    uploaded_file = st.file_uploader("Or upload a code file", type=["py", "js", "ts", "java", "cpp", "cs"])
    if uploaded_file:
        code_input = uploaded_file.read().decode("utf-8")
        st.text_area("File content", code_input, height=200, key="file_content")
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        programming_language = st.selectbox("Programming Language", PROGRAMMING_LANGUAGES)
    with col2:
        skill_level = st.selectbox("Skill Level", SKILL_LEVELS)
    with col3:
        user_role = st.selectbox("Your Role", USER_ROLES)
    with col4:
        explanation_language = st.selectbox("Explanation Language", EXPLANATION_LANGUAGES)
    if code_input:
        st.caption(f"Complexity: {calculate_code_complexity(code_input)}")
    if st.button("Run Workflow", type="primary"):
        if not code_input.strip():
            st.error("Please paste or upload your code.")
        elif not code_matches_language(code_input, programming_language):
            st.error(f"Language mismatch. Please check your code and language selection.")
        else:
            with st.spinner("Running Code Workflow..."):
                steps = [
                    ("Explain", call_groq_api(f"Explain this {programming_language} code for a {skill_level} {user_role} in {explanation_language}:\n{code_input}")),
                    ("Refactor", call_blackbox_agent([
    {"role": "system", "content": "You are a helpful coding assistant."},
    {"role": "user", "content": f"Refactor this {programming_language} code: {code_input}"}
])),
                    ("Review", call_groq_api(f"Review this {programming_language} code for errors and improvements: {code_input}")),
                    ("ErrorDetection", call_groq_api(f"Find bugs in this {programming_language} code: {code_input}")),
                    ("TestGeneration", call_groq_api(f"Generate tests for this {programming_language} code: {code_input}")),
                ]
                timeline = []
                for step, output in steps:
                    timeline.append({"step": step, "output": output})
                st.success("Workflow complete!")
                for t in timeline:
                    st.subheader(t["step"])
                    st.write(t["output"])
                # Show code diff (Original vs Refactored)
                st.subheader("Code Diff (Original vs Refactored)")
                refactored_code = steps[1][1]  # Blackbox agent output
                st.code(get_inline_diff(code_input, refactored_code), language=programming_language.lower())
                # Download report
                report = f"Code Workflow Report\nGenerated on: {datetime.datetime.now()}\nLanguage: {programming_language}\nSkill Level: {skill_level}\nRole: {user_role}\n\n"
                for t in timeline:
                    report += f"## {t['step']}\n{t['output']}\n\n---\n\n"
                st.download_button("Download Report", report, file_name="ai_workflow_report.txt")

elif page == "Semantic Search":
    st.header("Semantic Search")
    code_input = st.text_area("Paste your code here", height=200, key="sem_code")
    uploaded_file = st.file_uploader("Or upload a code file", type=["py", "js", "ts", "java", "cpp", "cs"], key="sem_file")
    if uploaded_file:
        code_input = uploaded_file.read().decode("utf-8")
        st.text_area("File content", code_input, height=200, key="sem_file_content")
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        programming_language = st.selectbox("Programming Language", PROGRAMMING_LANGUAGES, key="sem_lang")
    with col2:
        skill_level = st.selectbox("Skill Level", SKILL_LEVELS, key="sem_skill")
    with col3:
        user_role = st.selectbox("Your Role", USER_ROLES, key="sem_role")
    with col4:
        explanation_language = st.selectbox("Explanation Language", EXPLANATION_LANGUAGES, key="sem_expl")
    # Initialize session state variables for voice input and auto run
    if "voice_question" not in st.session_state:
        st.session_state.voice_question = ""
    if "auto_run_search" not in st.session_state:
        st.session_state.auto_run_search = False

    # Container for question input and voice button
    col_question, col_voice = st.columns([8,1])
    with col_question:
        question = st.text_input("Ask a question about your code", value=st.session_state.voice_question, key="question_input")
    with col_voice:
        # Microphone button with custom HTML and JS for voice input
        st.markdown(
            """
            <button id="mic-btn" title="Click to speak" style="height:38px; width:38px; font-size:20px;">🎤</button>
            <script>
            const micBtn = window.parent.document.querySelector('#mic-btn');
            const streamlitDoc = window.parent.document;

            // Use Web Speech API for voice recognition
            const SpeechRecognition = window.SpeechRecognition || window.webkitSpeechRecognition;
            if (SpeechRecognition) {
                const recognition = new SpeechRecognition();
                recognition.lang = 'en-US';
                recognition.interimResults = false;
                recognition.maxAlternatives = 1;

                micBtn.onclick = () => {
                    recognition.start();
                    micBtn.textContent = '🎙️';
                };

                recognition.onresult = (event) => {
                    const transcript = event.results[0][0].transcript;
                    // Send transcript to Streamlit via custom event
                    const inputEvent = new CustomEvent("voiceInput", {detail: transcript});
                    streamlitDoc.dispatchEvent(inputEvent);
                    micBtn.textContent = '🎤';
                };

                recognition.onerror = (event) => {
                    console.error('Speech recognition error', event.error);
                    micBtn.textContent = '🎤';
                };
            } else {
                micBtn.disabled = true;
                micBtn.title = "Speech Recognition not supported in this browser.";
            }
            </script>
            """,
            unsafe_allow_html=True
        )

    # Listen for the custom event and update session state via Streamlit's experimental_rerun hack
    # This requires a small hack using st.experimental_get_query_params and st.experimental_set_query_params
    # We will use st.experimental_get_query_params to detect voice input from URL params

    # Check if voice input is passed via query params
    query_params = st.experimental_get_query_params()
    if "voice_input" in query_params:
        voice_text = query_params["voice_input"][0]
        if voice_text != st.session_state.voice_question:
            st.session_state.voice_question = voice_text
            st.session_state.auto_run_search = True
            # Clear the query param to avoid repeated triggers
            st.experimental_set_query_params()

    # Run semantic search automatically if flag is set
    if st.session_state.auto_run_search:
        st.session_state.auto_run_search = False
        if not code_input.strip() or not st.session_state.voice_question.strip():
            st.error("Both code and question are required.")
        elif not code_matches_language(code_input, programming_language):
            st.error(f"Language mismatch. Please check your code and language selection.")
        else:
            with st.spinner("Running Semantic Search..."):
                answer = call_groq_api(f"{st.session_state.voice_question}\n\nCode:\n{code_input}")
                st.success("Answer:")
                st.write(answer)

    # Also keep the manual button for fallback
    if st.button("Run Semantic Search"):
        if not code_input.strip() or not question.strip():
            st.error("Both code and question are required.")
        elif not code_matches_language(code_input, programming_language):
            st.error(f"Language mismatch. Please check your code and language selection.")
        else:
            with st.spinner("Running Semantic Search..."):
                answer = call_groq_api(f"{question}\n\nCode:\n{code_input}")
                st.success("Answer:")
                st.write(answer)