import difflib
import streamlit as st
from groq import Groq
import os

# --- Set page config FIRST! ---
st.set_page_config(page_title="AI Code Assistant", layout="wide")

# --- Custom CSS for Professional Look ---
st.markdown("""
    <style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600&display=swap');
    html, body, [class*="css"]  {
        font-family: 'Inter', sans-serif;
        background-color: #f7f9fb;
    }
    .stApp {
        background-color: #f7f9fb;
    }
    .stSidebar {
        background-color: #22304a !important;
    }
    .stButton>button {
        background-color: #22304a;
        color: #fff;
        border-radius: 6px;
        border: none;
        font-weight: 600;
        padding: 0.5rem 1.5rem;
        margin-top: 0.5rem;
        margin-bottom: 0.5rem;
        transition: background 0.2s;
    }
    .stButton>button:hover {
        background-color: #1a2333;
        color: #fff;
    }
    .stTextInput>div>div>input, .stTextArea>div>textarea {
        background: #fff;
        border: 1px solid #d1d5db;
        border-radius: 6px;
        color: #22304a;
        font-size: 1rem;
    }
    .stDownloadButton>button {
        background-color: #22304a;
        color: #fff;
        border-radius: 6px;
        border: none;
        font-weight: 600;
        padding: 0.5rem 1.5rem;
        margin-top: 0.5rem;
        margin-bottom: 0.5rem;
        transition: background 0.2s;
    }
    .stDownloadButton>button:hover {
        background-color: #1a2333;
        color: #fff;
    }
    .stExpanderHeader {
        font-weight: 600;
        color: #22304a;
        font-size: 1.1rem;
    }
    .stMarkdown {
        color: #22304a;
    }
    .stAlert {
        border-radius: 6px;
    }
    </style>
""", unsafe_allow_html=True)

# --- Groq API Setup ---
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
if not GROQ_API_KEY:
    st.error("GROQ_API_KEY environment variable not set. Please set it in your Hugging Face Space secrets.")
    st.stop()
client = Groq(api_key=GROQ_API_KEY)

def groq_api_call(prompt):
    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama3-70b-8192",
    )
    return chat_completion.choices[0].message.content

def get_diff_html(original, modified):
    original_lines = original.splitlines()
    modified_lines = modified.splitlines()
    differ = difflib.HtmlDiff(tabsize=4, wrapcolumn=80)
    return differ.make_table(original_lines, modified_lines, "Original", "Modified", context=True, numlines=2)

def code_complexity(code):
    lines = code.count('\n') + 1
    functions = code.count('def ')
    classes = code.count('class ')
    comments = code.count('#')
    return f"Lines: {lines}, Functions: {functions}, Classes: {classes}, Comments: {comments}"

def detect_code_type(code, programming_language):
    backend_keywords = [
        'flask', 'django', 'express', 'fastapi', 'spring', 'controller', 'api', 'server', 'database', 'sql', 'mongoose'
    ]
    frontend_keywords = [
        'react', 'vue', 'angular', 'component', 'html', 'css', 'document.getelementbyid', 'window.', 'render', 'jsx',
        '<html', '<body', '<script', '<div', 'getelementbyid', 'queryselector', 'addeventlistener', 'innerhtml'
    ]
    data_science_keywords = [
        'pandas', 'numpy', 'sklearn', 'matplotlib', 'seaborn', 'plt', 'train_test_split', 'randomforestclassifier', 'classification_report'
    ]
    code_lower = code.lower()
    if any(word in code_lower for word in data_science_keywords):
        return 'data_science'
    if any(word in code_lower for word in frontend_keywords):
        return 'frontend'
    if programming_language.lower() in ['python', 'java', 'c#']:
        if any(word in code_lower for word in backend_keywords):
            return 'backend'
    if programming_language.lower() in ['javascript', 'typescript', 'java', 'c#']:
        if any(word in code_lower for word in frontend_keywords):
            return 'frontend'
    if programming_language.lower() in ['python', 'java', 'c#']:
        return 'backend'
    if programming_language.lower() in ['javascript', 'typescript']:
        return 'frontend'
    return 'unknown'

def code_matches_language(code: str, language: str) -> bool:
    code = code.strip().lower()
    if language.lower() == "python":
        return "def " in code or "import " in code or ".py" in code
    if language.lower() == "c++":
        return "#include" in code or "int main" in code or ".cpp" in code or "std::" in code
    if language.lower() == "java":
        return "public class" in code or "public static void main" in code or ".java" in code
    if language.lower() == "c#":
        return "using system" in code or "namespace" in code or ".cs" in code
    if language.lower() == "javascript":
        return "function " in code or "const " in code or "let " in code or "var " in code or ".js" in code
    if language.lower() == "typescript":
        return "function " in code or "const " in code or "let " in code or "var " in code or ": string" in code or ".ts" in code
    if language.lower() == "html":
        return "<html" in code or "<!doctype html" in code
    return True  # fallback

def agentic_workflow(code, skill_level, programming_language, explanation_language, user_role):
    timeline = []
    suggestions = []

    explain_prompt = (
        f"Explain the following {programming_language} code line by line or function by function "
        f"for a {skill_level.lower()} {user_role} in {explanation_language}:\n{code}"
    )
    explanation = groq_api_call(explain_prompt)
    timeline.append({
        "step": "Explain",
        "description": "Step-by-step explanation of your code.",
        "output": explanation,
        "code": code
    })
    suggestions.append("Refactor your code for better readability and performance.")

    refactor_prompt = (
        f"Refactor the following {programming_language} code for better readability, performance, and structure. "
        f"Explain what changes you made and why, for a {skill_level.lower()} {user_role} in {explanation_language}:\n{code}"
    )
    refactor_response = groq_api_call(refactor_prompt)
    if "```" in refactor_response:
        parts = refactor_response.split("```")
        refactor_explanation = parts[0].strip()
        refactored_code = ""
        for i in range(1, len(parts)):
            if parts[i].strip().startswith(programming_language.lower()):
                refactored_code = parts[i].strip().split('\n', 1)[1] if '\n' in parts[i] else ""
                break
            elif i == 1:
                refactored_code = parts[i].strip().split('\n', 1)[1] if '\n' in parts[i] else ""
        if not refactored_code:
            refactored_code = refactor_response.strip()
    else:
        refactor_explanation = "Refactored code below."
        refactored_code = refactor_response.strip()
    timeline.append({
        "step": "Refactor",
        "description": refactor_explanation,
        "output": refactored_code,
        "code": refactored_code
    })
    suggestions.append("Review the refactored code for best practices and improvements.")

    review_prompt = (
        f"Provide a code review for the following {programming_language} code. "
        f"Include feedback on best practices, code smells, optimization, and security issues, for a {skill_level.lower()} {user_role} in {explanation_language}:\n{refactored_code}"
    )
    review_feedback = groq_api_call(review_prompt)
    timeline.append({
        "step": "Review",
        "description": "AI code review and feedback.",
        "output": review_feedback,
        "code": refactored_code
    })
    suggestions.append("Generate unit tests for your code.")

    test_prompt = (
        f"Write unit tests for the following {programming_language} code. "
        f"Use pytest style and cover all functions. For a {skill_level.lower()} {user_role} in {explanation_language}:\n{refactored_code}"
    )
    test_code = groq_api_call(test_prompt)
    timeline.append({
        "step": "Test Generation",
        "description": "AI-generated unit tests for your code.",
        "output": test_code,
        "code": test_code
    })
    suggestions.append("Run the generated tests in your local environment.")

    return timeline, suggestions

st.markdown(
    "<h2 style='text-align: center; color: #22304a; font-weight: 600; margin-bottom: 0.5em;'>AI Code Assistant</h2>",
    unsafe_allow_html=True
)

with st.sidebar:
    st.title("Settings")
    programming_language = st.selectbox(
        "Programming Language",
        ["Python", "C++", "Java", "C#", "JavaScript", "TypeScript", "HTML"]
    )
    explanation_language = st.selectbox(
        "Explanation Language",
        ["English", "Urdu", "Chinese", "Spanish"]
    )
    skill_level = st.selectbox("Skill Level", ["Beginner", "Intermediate", "Expert"])
    user_role = st.selectbox(
        "Choose Role",
        ["Data Scientist", "Backend Developer", "Frontend Developer", "Student"]
    )
    st.markdown("---")
    st.markdown("<span style='color:#fff;'>Powered by <b>BLACKBOX.AI</b></span>", unsafe_allow_html=True)

if "code" not in st.session_state:
    st.session_state.code = ""
if "timeline" not in st.session_state:
    st.session_state.timeline = []
if "suggestions" not in st.session_state:
    st.session_state.suggestions = []

col1, col2 = st.columns([2, 3], gap="large")

with col1:
    st.subheader(f"{programming_language} Code")
    uploaded_file = st.file_uploader(f"Upload .{programming_language.lower()} file", type=[programming_language.lower()])
    code_input = st.text_area(
        "Paste or edit your code here:",
        height=300,
        value=st.session_state.code,
        key="main_code_input"
    )
    if uploaded_file is not None:
        code = uploaded_file.read().decode("utf-8")
        st.session_state.code = code
        st.success("File uploaded successfully.")
    elif code_input:
        st.session_state.code = code_input

    st.markdown(f"<b>Complexity:</b> {code_complexity(st.session_state.code)}", unsafe_allow_html=True)

    st.markdown("---")
    st.markdown("#### Agent Suggestions")
    for suggestion in st.session_state.suggestions[-3:]:
        st.markdown(f"- {suggestion}")

    st.markdown("---")
    st.markdown("#### Download Full Report")
    if st.session_state.timeline:
        report = ""
        for step in st.session_state.timeline:
            report += f"## {step['step']}\n{step['description']}\n\n{step['output']}\n\n"
        st.download_button("Download Report", report, file_name="ai_code_assistant_report.txt")

with col2:
    st.subheader("Agentic Workflow")
    if st.button("Run Full AI Agent Workflow"):
        if not st.session_state.code.strip():
            st.warning("Please enter or upload code first.")
        else:
            # Language check
            if not code_matches_language(st.session_state.code, programming_language):
                st.error(f"It looks like you provided code in a different language. Please provide {programming_language} code.")
            else:
                code_type = detect_code_type(st.session_state.code, programming_language)
                # Role/code type enforcement
                if code_type == "data_science" and user_role != "Data Scientist":
                    st.error("It looks like you provided data science code. Please select 'Data Scientist' as your role.")
                elif code_type == "frontend" and user_role != "Frontend Developer":
                    st.error("It looks like you provided frontend code. Please select 'Frontend Developer' as your role.")
                elif code_type == "backend" and user_role != "Backend Developer":
                    st.error("It looks like you provided backend code. Please select 'Backend Developer' as your role.")
                elif code_type == "unknown":
                    st.warning("Could not determine the code type. Please make sure your code is complete and clear.")
                else:
                    with st.spinner("AI Agent is working through all steps..."):
                        timeline, suggestions = agentic_workflow(
                            st.session_state.code,
                            skill_level,
                            programming_language,
                            explanation_language,
                            user_role
                        )
                        st.session_state.timeline = timeline
                        st.session_state.suggestions = suggestions
                    st.success("Agentic workflow complete. See timeline below.")

    if st.session_state.timeline:
        for i, step in enumerate(st.session_state.timeline):
            with st.expander(f"Step {i+1}: {step['step']}"):
                st.markdown(f"<b>{step['description']}</b>", unsafe_allow_html=True)
                if step['step'] in ["Refactor", "Test Generation"]:
                    if i > 0:
                        prev_code = st.session_state.timeline[i-1]['code']
                        diff_html = get_diff_html(prev_code, step['code'])
                        st.markdown("**Side-by-Side Diff:**")
                        st.components.v1.html(diff_html, height=400, scrolling=True)
                if step['step'] in ["Refactor", "Test Generation"]:
                    st.markdown("**Code Output:**")
                    st.code(step['output'], language=programming_language.lower())
                else:
                    st.markdown("**Output:**")
                    st.write(step['output'])
    else:
        st.info("Run the agentic workflow to see step-by-step results, explanations, and code evolution.")

st.markdown("---")
st.markdown('<div style="text-align: center; color: #22304a; font-size: 1rem; margin-top: 2em;">Powered by <b>BLACKBOX.AI</b></div>', unsafe_allow_html=True)