import streamlit as st
import os
import json
import pandas as pd
from dotenv import load_dotenv
load_dotenv()
from data.bug_dataset import TRAINING_SCENARIOS
from orchestrator import Debugger

st.set_page_config(page_title="CodeDebugger RL", layout="wide")

st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
    
    html, body, [class*="css"]  {
        font-family: 'Inter', sans-serif;
    }
    
    /* Background and typography */
    .stApp {
        background-color: #0f111a;
        color: #e2e8f0;
    }
    
    /* Headers with gradient */
    h1, h2, h3 {
        background: linear-gradient(90deg, #3b82f6, #8b5cf6, #ec4899);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-weight: 700 !important;
    }
    
    /* Elegant card styling for blocks */
    .stCodeBlock, div[data-testid="stText"] {
        background: rgba(30, 41, 59, 0.5) !important;
        border: 1px solid rgba(255, 255, 255, 0.1);
        border-radius: 12px;
        padding: 4px;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
    }
    
    /* Beautiful modern buttons */
    .stButton>button {
        background: linear-gradient(90deg, #3b82f6, #8b5cf6);
        color: white;
        border: none;
        border-radius: 8px;
        padding: 0.5rem 1rem;
        font-weight: 600;
        transition: all 0.3s ease;
        box-shadow: 0 4px 15px rgba(139, 92, 246, 0.3);
    }
    
    .stButton>button:hover {
        transform: translateY(-2px);
        box-shadow: 0 6px 20px rgba(139, 92, 246, 0.5);
        color: white;
    }
    
    /* Input fields */
    .stTextInput>div>div>input, .stTextArea>div>textarea, .stSelectbox>div>div {
        background-color: rgba(15, 23, 42, 0.8);
        border: 1px solid rgba(255, 255, 255, 0.2);
        color: white;
        border-radius: 8px;
    }
    
    .stTextInput>div>div>input:focus, .stTextArea>div>textarea:focus {
        border-color: #8b5cf6;
        box-shadow: 0 0 0 1px #8b5cf6;
    }
    
    /* Welcome screen empty state */
    .empty-state {
        display: flex;
        flex-direction: column;
        align-items: center;
        justify-content: center;
        padding: 60px 20px;
        text-align: center;
        background: rgba(30, 41, 59, 0.3);
        border: 1px dashed rgba(255, 255, 255, 0.2);
        border-radius: 16px;
        margin-top: 40px;
    }
    .empty-state h3 {
        margin-bottom: 10px;
        background: none;
        -webkit-text-fill-color: #e2e8f0;
        color: #e2e8f0;
    }
    .empty-state p {
        color: #94a3b8;
        max-width: 500px;
        margin: 0 auto;
        line-height: 1.6;
    }
</style>
""", unsafe_allow_html=True)

# Sync API Key from state to env
if "groq_api_key" not in st.session_state:
    st.session_state.groq_api_key = os.environ.get("GROQ_API_KEY", "")

# Load datasets
problems_dict = {f"{p['id']}: {p.get('title', 'Unknown')} [{p['difficulty']}]": p for p in TRAINING_SCENARIOS}

st.title("CodeDebugger RL Environment")

tab1, tab2, tab3 = st.tabs(["Debug Results", "Before vs After Training", "Training Results"])

# ==========================================
# TAB 1: Debug Results
# ==========================================
with tab1:
    with st.sidebar:
        st.header("Debugger Settings")
        
        # Select Problem
        custom_toggle = st.checkbox("Paste custom buggy code")
        if custom_toggle:
            custom_code = st.text_area("Buggy Code")
            custom_error = st.text_input("Error Type (e.g. IndexError)")
            problem = {
                "id": "custom",
                "difficulty": "custom",
                "title": "Custom Code",
                "description": "Custom bug",
                "buggy_code": custom_code,
                "error_type": custom_error,
                "test_cases": [{"input": "N/A", "expected": "N/A"}] # dummy
            }
        else:
            selected_prob_name = st.selectbox("Select Problem", list(problems_dict.keys()))
            problem = problems_dict[selected_prob_name]
            
        iterations = st.slider("Max Iterations", 1, 5, 3)
        
        api_key = st.text_input("Groq API Key (Optional)", value=st.session_state.groq_api_key, type="password")
        if api_key:
            os.environ["GROQ_API_KEY"] = api_key
            st.session_state.groq_api_key = api_key
            
        run_btn = st.button("Run Debugger", type="primary")

    if run_btn:
        with st.spinner("Running Debugger..."):
            debugger = Debugger(max_iterations=iterations)
            result = debugger.run(problem, verbose=False)
            st.session_state.last_result = result
            
    if "last_result" in st.session_state:
        res = st.session_state.last_result
        
        st.header(f"Best Fix: Iter {res['best_iter']} | Reward: {res['best_reward']:.2f}")
        
        # Action Buttons
        col1, col2 = st.columns(2)
        with col1:
            if st.button("View Iteration History", use_container_width=True):
                st.session_state.show_history = True
        with col2:
            if st.button("View Reward Details", use_container_width=True):
                st.session_state.show_reward = True
                
        # Main Area Code
        st.markdown("### Best Fixed Code")
        st.code(res.get("best_code", "No code returned"), language="python")
        
        # Test Results & Anti-Hack
        passed = res['tests_passed_final']
        total = res['tests_total']
        
        st.markdown(f"**Test Results:** {passed}/{total} Passed")
        
        last_iter_info = res["iterations"][-1]
        safety = last_iter_info["info"].get("safety_violations", [])
        anti_hack_failed = last_iter_info["info"].get("anti_hack_failed", False)
        
        if not safety and not anti_hack_failed:
            st.success("✅ No cheating detected")
        else:
            violations = safety.copy() if safety else []
            if anti_hack_failed: violations.append("Anti-hack checks failed")
            st.error(f"🚨 Cheating caught: {', '.join(violations)}")
    else:
        st.markdown("""
        <div class="empty-state">
            <div style="font-size: 48px; margin-bottom: 15px;">✨</div>
            <h3>Ready to Debug</h3>
            <p>Select a buggy Python problem from the sidebar or paste your own custom code. Click <b>Run Debugger</b> to unleash the RL agent to fix your code, score it, and ensure no anti-hacking rules are broken.</p>
        </div>
        """, unsafe_allow_html=True)

# Modals for Tab 1
@st.dialog("Iteration History", width="large")
def show_history_modal():
    res = st.session_state.last_result
    for i, it in enumerate(res["iterations"]):
        r_total = it["reward"]["total"]
        if r_total >= 80: icon = "✅"
        elif r_total >= 50: icon = "⚠️"
        else: icon = "❌"
            
        with st.expander(f"Iteration {it['iteration']} | Reward: {r_total:.2f} {icon}"):
            method = it["fix_result"].get("method", "unknown")
            st.caption(f"Method: {method}")
            
            st.markdown("#### Component Scores")
            html = ""
            for k, v in it["reward"].items():
                if k != "total":
                    color = "green" if v >= 0 else "red"
                    # Simple heuristic scaling for visual bar
                    perc = min(100, max(0, abs(v) * 2)) 
                    html += f"""
                    <div style='margin-bottom:4px; font-size: 14px;'>
                        {k} ({v:.1f})
                        <div style='background-color:#e0e0e0; width:100%; height:10px; border-radius:5px;'>
                            <div style='background-color:{color}; width:{perc}%; height:10px; border-radius:5px;'></div>
                        </div>
                    </div>
                    """
            st.markdown(html, unsafe_allow_html=True)
            
            st.markdown("#### Code Submitted")
            st.code(it["fix_result"].get("fixed_code", ""), language="python")

@st.dialog("Reward Details", width="large")
def show_reward_modal():
    res = st.session_state.last_result
    
    st.markdown("### Reward Trajectory")
    iters = [it["iteration"] for it in res["iterations"]]
    rewards = [it["reward"]["total"] for it in res["iterations"]]
    df = pd.DataFrame({"Iteration": iters, "Reward": rewards}).set_index("Iteration")
    st.line_chart(df)
    
    st.markdown("### Component Breakdown (Final Iteration)")
    final_reward_dict = res["iterations"][-1]["reward"]
    components = {k: v for k, v in final_reward_dict.items() if k != "total"}
    
    # Horizontal bar chart
    comp_df = pd.DataFrame(list(components.items()), columns=["Component", "Score"]).set_index("Component")
    st.bar_chart(comp_df, horizontal=True)

# Trigger dialogs based on session state
if st.session_state.get("show_history", False):
    show_history_modal()
    st.session_state.show_history = False

if st.session_state.get("show_reward", False):
    show_reward_modal()
    st.session_state.show_reward = False

# ==========================================
# TAB 2: Before vs After Training
# ==========================================
with tab2:
    b_path = "outputs/baseline_scores.json"
    t_path = "outputs/trained_scores.json"
    
    if not os.path.exists(b_path) or not os.path.exists(t_path):
        st.info("Run run_baseline.py and train_grpo.py first to generate outputs/baseline_scores.json and trained_scores.json")
    else:
        try:
            with open(b_path, "r") as f: b_data = json.load(f)
            with open(t_path, "r") as f: t_data = json.load(f)
            
            baseline = b_data.get("results", b_data) if isinstance(b_data, dict) else b_data
            trained = t_data.get("results", t_data) if isinstance(t_data, dict) else t_data
                
            base_solved = sum(1 for p in baseline if p.get("solved"))
            train_solved = sum(1 for p in trained if p.get("solved"))
            total = len(baseline)
            
            st.header(f"Solved: {base_solved}/{total} → {train_solved}/{total} after training")
            
            prob_ids = [p["problem_id"] for p in baseline]
            df_chart = pd.DataFrame({
                "Baseline": [p["best_reward"] for p in baseline],
                "Trained": [p["best_reward"] for p in trained]
            }, index=prob_ids)
            
            st.bar_chart(df_chart)
            
            # Difficulty Table
            st.markdown("### Difficulty Breakdown")
            diff_metrics = []
            for diff in ["easy", "medium", "hard"]:
                b_diff = [p for p in baseline if p.get("difficulty") == diff]
                t_diff = [p for p in trained if p.get("difficulty") == diff]
                if not b_diff: continue
                
                b_avg = sum(p["best_reward"] for p in b_diff) / len(b_diff)
                t_avg = sum(p["best_reward"] for p in t_diff) / len(t_diff)
                
                diff_metrics.append({
                    "Difficulty": diff.capitalize(),
                    "Before (Avg)": round(b_avg, 2),
                    "After (Avg)": round(t_avg, 2),
                    "Delta": round(t_avg - b_avg, 2)
                })
            st.dataframe(pd.DataFrame(diff_metrics))
                
        except Exception as e:
            st.error(f"Error parsing JSON files: {e}")

# ==========================================
# TAB 3: Training Results
# ==========================================
with tab3:
    st.header("Training Results")
    log_path = "outputs/training_log.jsonl"
    
    if not os.path.exists(log_path):
        st.info("Run train_grpo.py first to generate outputs/training_log.jsonl")
    else:
        try:
            steps, rewards, diffs = [], [], []
            with open(log_path, "r") as f:
                for line in f:
                    if line.strip():
                        data = json.loads(line)
                        steps.append(data.get("step", len(steps)))
                        rewards.append(data.get("reward", 0.0))
                        diffs.append(data.get("difficulty", "unknown"))
                        
            df_train = pd.DataFrame({"Reward": rewards, "Difficulty": diffs}, index=steps)
            st.line_chart(df_train, y="Reward", color="Difficulty")
            
        except Exception as e:
            st.error(f"Error parsing training logs: {e}")