#!/usr/bin/env python3

import os
import sys
import logging

import uvicorn
import gradio as gr
from typing import Dict, Any, Optional, List

from fastapi import Body, FastAPI
from pydantic import BaseModel, ConfigDict, Field

# Add current directory to Python path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from env.environment import CodeGuardianEnv
from env.models import Action, Observation
import plotly.express as px
import pandas as pd

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class CodeGuardianSpace:
    def __init__(self):
        self.env = CodeGuardianEnv()
        self.current_obs = None
        self.baseline_score = None
        self.rl_best_score = None
        self.episode_history = []
        self.reset_environment()
    
    def reset_environment(self):
        """Reset the environment and return initial observation."""
        try:
            self.current_obs = self.env.reset()
            return self.format_observation(self.current_obs)
        except Exception as e:
            logger.error(f"Failed to reset environment: {e}")
            return f"Error: {str(e)}"
    
    def format_observation(self, obs: Observation) -> str:
        """Format observation for display."""
        return f"""
**Current Task:** {obs.task_id} ({obs.difficulty})
**Step:** {obs.step_index}/3

**Code:**
```python
{obs.code}
```

**Intent:** {obs.intent}

**Constraints:**
{chr(10).join(f"- {c}" for c in obs.constraints)}
"""
    
    def step_environment(self, action_type: str, detected_issues: str = "", risk_level: str = "low", suggested_fix: str = ""):
        """Execute a step in the environment with enhanced feedback."""
        try:
            if not self.current_obs or self.env.done:
                return "Environment is done. Please reset to continue.", "", 0.0, "", "", ""
            
            # Parse detected issues
            issues_list = []
            if detected_issues.strip():
                issues_list = [issue.strip() for issue in detected_issues.split(',')]
            
            # Create action
            action = Action(
                action_type=action_type,
                detected_issues=issues_list if action_type == "detect" else None,
                risk_level=risk_level if action_type == "risk" else None,
                suggested_fix=suggested_fix if action_type == "fix" else None
            )
            
            # Execute step
            next_obs, reward, done = self.env.step(action)
            self.current_obs = next_obs
            
            # Format enhanced results
            obs_text = self.format_observation(next_obs) if next_obs else "Environment completed!"
            feedback_text = f"Feedback: {reward.feedback}" if reward.feedback else "No feedback"
            current_score = self.env.state()['partial_score']
            
            # Format reward breakdown
            breakdown_text = ""
            if reward.breakdown:
                breakdown_parts = []
                for step_type, score in reward.breakdown.items():
                    if score > 0:
                        breakdown_parts.append(f"{step_type}: {score:.3f}")
                breakdown_text = " | ".join(breakdown_parts) if breakdown_parts else "No breakdown available"
            
            # Format detailed step info
            step_details = f"""
**Step Details:**
- Action: {action.action_type}
- Reward: {reward.value:+.3f}
- Partial Scores: {reward.partial_scores}
- Breakdown: {breakdown_text}
""" if reward.breakdown else ""
            
            score_text = f"Current Score: {current_score:.2f}"
            
            if done:
                obs_text += f"\n\n**🏁 Final Score: {current_score:.2f}**"
                obs_text += f"\n**Total Steps: {len(self.env.state())}**"
            
            return obs_text, feedback_text, reward.value, breakdown_text, step_details
            
        except Exception as e:
            logger.error(f"Failed to execute step: {e}")
            return f"Error: {str(e)}", "", 0.0, "", ""
    
    def run_full_episode(self) -> str:
        """Run a complete episode automatically and return results."""
        try:
            self.reset_environment()
            episode_log = []
            total_reward = 0.0
            step_count = 0
            
            while not self.env.done and step_count < 50:  # Safety limit
                current_step = self.env.state()['current_step']
                step_types = ['detect', 'risk', 'fix']
                step_type = step_types[current_step]
                
                # Create mock action (in real usage, this would be from LLM)
                if step_type == 'detect':
                    action = Action(action_type='detect', detected_issues=['mock issue'])
                elif step_type == 'risk':
                    action = Action(action_type='risk', risk_level='medium')
                else:
                    action = Action(action_type='fix', suggested_fix='# Mock fix')
                
                next_obs, reward, done = self.env.step(action)
                self.current_obs = next_obs
                total_reward += reward.value
                step_count += 1
                
                episode_log.append(f"Step {step_count}: {step_type} -> {reward.value:+.3f}")
            
            # Add to episode history
            self.episode_history.append(total_reward)
            
            result = f"""
**🎯 Episode Complete!**

**Summary:**
- Total Steps: {step_count}
- Final Score: {total_reward:.3f}
- Tasks Completed: {self.env.state()['current_task']}

**Step-by-Step:**
{chr(10).join(episode_log)}
"""
            
            return result
            
        except Exception as e:
            logger.error(f"Failed to run full episode: {e}")
            return f"Error: {str(e)}"
    
    def run_baseline_comparison(self) -> str:
        """Run baseline agent for comparison."""
        try:
            # Run baseline
            baseline_env = CodeGuardianEnv()
            obs = baseline_env.reset()
            baseline_reward = 0.0
            step_count = 0
            
            while not baseline_env.done and step_count < 50:
                current_step = baseline_env.state()['current_step']
                step_types = ['detect', 'risk', 'fix']
                step_type = step_types[current_step]
                
                # Baseline always outputs minimal responses
                if step_type == 'detect':
                    action = Action(action_type='detect', detected_issues=[])
                elif step_type == 'risk':
                    action = Action(action_type='risk', risk_level='low')
                else:
                    action = Action(action_type='fix', suggested_fix='')
                
                next_obs, reward, done = baseline_env.step(action)
                baseline_reward += reward.value
                step_count += 1
                obs = next_obs
            
            self.baseline_score = baseline_reward
            
            return f"""
**📊 Baseline Complete!**

**Baseline Results:**
- Total Steps: {step_count}
- Baseline Score: {baseline_reward:.3f}
- Strategy: Minimal (empty responses)
"""
            
        except Exception as e:
            logger.error(f"Failed to run baseline: {e}")
            return f"Error: {str(e)}"
    
    def get_performance_summary(self) -> str:
        """Get visual performance summary with improvement calculation."""
        if self.baseline_score is None:
            return "Run baseline comparison first to see improvement metrics."
        
        if not self.episode_history:
            return "Run episodes to see performance data."
        
        best_score = max(self.episode_history)
        latest_score = self.episode_history[-1]
        
        # Calculate improvement
        improvement = ((best_score - self.baseline_score) / (abs(self.baseline_score) + 1e-6)) * 100
        
        # Performance classification
        if improvement > 50:
            emoji = "🚀"
            status = "EXCELLENT"
        elif improvement > 20:
            emoji = "✅"
            status = "GOOD"
        elif improvement > 0:
            emoji = "📈"
            status = "POSITIVE"
        else:
            emoji = "⚠️"
            status = "NO IMPROVEMENT"
        
        return f"""
**{emoji} PERFORMANCE COMPARISON**

**Baseline Score:** {self.baseline_score:.3f}
**Best RL Score:** {best_score:.3f}
**Latest Score:** {latest_score:.3f}
**Improvement:** +{improvement:.1f}%
**Status:** {status}

**Episode Progression:**
{chr(10).join(f'Episode {i+1}: {score:.3f}' for i, score in enumerate(self.episode_history))}
"""
    
    def create_progress_plot(self):
        """Create episode progression plot."""
        if not self.episode_history:
            return None
        
        df = pd.DataFrame({
            'Episode': range(1, len(self.episode_history) + 1),
            'Score': self.episode_history
        })
        
        fig = px.line(
            df, 
            x='Episode', 
            y='Score',
            title='Episode Score Progression',
            markers=True,
            line_shape='linear'
        )
        
        # Add baseline line if available
        if self.baseline_score is not None:
            fig.add_hline(
                y=self.baseline_score, 
                line_dash="dash", 
                line_color="red",
                annotation_text=f"Baseline: {self.baseline_score:.3f}",
                annotation_position="top right"
            )
        
        fig.update_layout(
            xaxis_title="Episode",
            yaxis_title="Score",
            height=400
        )
        
        return fig
    
    def get_expected_step(self):
        """Get the expected action type for current step."""
        if not self.current_obs or self.env.done:
            return "reset"
        step_types = ["detect", "risk", "fix"]
        return step_types[self.env.state()['current_step']]

# Initialize the space
space = CodeGuardianSpace()

def reset_interface():
    """Reset the environment interface."""
    return space.reset_environment()

def step_interface(action_type, detected_issues, risk_level, suggested_fix):
    """Step the environment interface."""
    return space.step_environment(action_type, detected_issues, risk_level, suggested_fix)

def run_episode_interface():
    """Run full episode interface."""
    return space.run_full_episode()

def run_baseline_interface():
    """Run baseline comparison interface."""
    return space.run_baseline_comparison()

def get_performance_interface():
    """Get performance summary interface."""
    return space.get_performance_summary()

def get_plot_interface():
    """Get progress plot interface."""
    return space.create_progress_plot()

def update_ui():
    """Update UI based on current state."""
    expected = space.get_expected_step()
    if expected == "reset":
        return (
            gr.update(visible=True), 
            gr.update(visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False),
            gr.update(visible=False)
        )
    else:
        return (
            gr.update(visible=False), 
            gr.update(visible=True), 
            gr.update(visible=True), 
            gr.update(visible=True), 
            gr.update(visible=True),
            gr.update(visible=True)
        )

# OpenEnv / TRL-style HTTP API — register routes on FastAPI before mounting Gradio at "/".


class OpenEnvStepRequest(BaseModel):
    model_config = ConfigDict(extra="allow")

    action: Dict[str, Any]
    timeout_s: Optional[float] = None
    request_id: Optional[str] = None


class OpenEnvStepResponse(BaseModel):
    model_config = ConfigDict(extra="forbid")

    observation: Dict[str, Any] = Field(default_factory=dict)
    reward: Optional[float] = None
    done: bool = False


app = FastAPI()


@app.post("/reset")
def reset():
    try:
        obs = space.env.reset()
        space.current_obs = obs

        if hasattr(obs, "model_dump"):
            obs = obs.model_dump()

        return {
            "observation": obs,
            "reward": 0.0,
            "done": False,
        }
    except Exception as e:
        print("RESET ERROR:", e)
        return {
            "observation": {},
            "reward": 0.0,
            "done": False,
        }


@app.post("/step", response_model=OpenEnvStepResponse)
async def openenv_step(request: OpenEnvStepRequest) -> OpenEnvStepResponse:
    action = Action.model_validate(request.action)
    next_obs, reward_obj, done = space.env.step(action)
    space.current_obs = next_obs
    obs_dict = next_obs.model_dump() if next_obs is not None else {}
    return OpenEnvStepResponse(
        observation=obs_dict,
        reward=float(reward_obj.value),
        done=done,
    )


@app.get("/state")
async def openenv_state() -> Dict[str, Any]:
    return space.env.state()


# Gradio UI — built after API routes; mounted last so /reset, /step, /state stay reachable.
with gr.Blocks(title="CodeGuardian Environment - RL Enhanced") as demo:
    gr.Markdown("# 🤖 CodeGuardian Environment - RL Enhanced")
    gr.Markdown("Test your AI agent on code review tasks: detect issues, assess risks, and suggest fixes with reinforcement learning.")

    with gr.Row():
        reset_btn = gr.Button("🔄 Reset Environment", variant="primary")
        baseline_btn = gr.Button("📊 Run Baseline", variant="secondary")
        episode_btn = gr.Button("🎯 Run Episode", variant="secondary")

    with gr.Row():
        perf_btn = gr.Button("📈 Show Performance", variant="secondary")
        plot_btn = gr.Button("📊 Show Progress Plot", variant="secondary")

    obs_display = gr.Markdown("Click 'Reset Environment' to start")

    with gr.Group(visible=True) as perf_group:
        gr.Markdown("## 📊 Performance Comparison")
        performance_display = gr.Markdown("Run baseline and episodes to see performance comparison.")

    with gr.Group(visible=True) as plot_group:
        gr.Markdown("## 📈 Episode Progression")
        progress_plot = gr.Plot(label="Score Progression")

    with gr.Group(visible=False) as action_group:
        gr.Markdown("## Take Action")

        with gr.Row():
            with gr.Column():
                action_type = gr.Radio(
                    choices=["detect", "risk", "fix"],
                    label="Action Type",
                    value="detect",
                )

                with gr.Group(visible=True) as detect_group:
                    detected_issues = gr.Textbox(
                        label="Detected Issues (comma-separated)",
                        placeholder="e.g., wrong operator, security issue, performance problem",
                    )

                with gr.Group(visible=False) as risk_group:
                    risk_level = gr.Radio(
                        choices=["low", "medium", "high"],
                        label="Risk Level",
                        value="low",
                    )

                with gr.Group(visible=False) as fix_group:
                    suggested_fix = gr.Textbox(
                        label="Suggested Fix",
                        lines=3,
                        placeholder="Describe or provide the fix...",
                    )

            with gr.Column():
                step_btn = gr.Button("⚡ Execute Step", variant="primary")
                feedback = gr.Textbox(label="📝 Feedback", interactive=False, lines=2)
                reward = gr.Number(label="🏆 Reward", interactive=False)

                reward_breakdown = gr.Textbox(label="📊 Reward Breakdown", interactive=False, lines=1)
                step_details = gr.Markdown("", label="Step Details")

    episode_results = gr.Markdown("", visible=False, label="Episode Results")

    reset_btn.click(
        reset_interface,
        outputs=obs_display,
    ).then(
        update_ui,
        outputs=[reset_btn, baseline_btn, episode_btn, perf_btn, plot_btn, action_group, detect_group, risk_group, fix_group],
    )

    baseline_btn.click(
        run_baseline_interface,
        outputs=performance_display,
    ).then(
        lambda: gr.update(visible=True),
        outputs=perf_group,
    )

    episode_btn.click(
        run_episode_interface,
        outputs=episode_results,
    ).then(
        lambda: gr.update(visible=True),
        outputs=episode_results,
    ).then(
        get_plot_interface,
        outputs=progress_plot,
    ).then(
        lambda: gr.update(visible=True),
        outputs=plot_group,
    )

    perf_btn.click(
        get_performance_interface,
        outputs=performance_display,
    ).then(
        lambda: gr.update(visible=True),
        outputs=perf_group,
    )

    plot_btn.click(
        get_plot_interface,
        outputs=progress_plot,
    ).then(
        lambda: gr.update(visible=True),
        outputs=plot_group,
    )

    def toggle_action_inputs(action):
        return (
            gr.update(visible=(action == "detect")),
            gr.update(visible=(action == "risk")),
            gr.update(visible=(action == "fix")),
        )

    action_type.change(
        toggle_action_inputs,
        inputs=action_type,
        outputs=[detect_group, risk_group, fix_group],
    )

    step_btn.click(
        step_interface,
        inputs=[action_type, detected_issues, risk_level, suggested_fix],
        outputs=[obs_display, feedback, reward, reward_breakdown, step_details],
    ).then(
        update_ui,
        outputs=[reset_btn, baseline_btn, episode_btn, perf_btn, plot_btn, action_group, detect_group, risk_group, fix_group],
    )

demo.queue()
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    print("SERVER STARTED ON PORT", port)
    uvicorn.run("app:app", host="0.0.0.0", port=port)