#!/usr/bin/env python3

import os
import json
import logging
from typing import Dict, Any, Optional, List

try:
    from openai import OpenAI
except ImportError:
    print("Error: openai package not found. Install with: pip install openai")
    exit(1)

from env.environment import CodeGuardianEnv
from env.models import Action, AgentResponse, AgentState

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def setup_openai_client() -> tuple[OpenAI, str]:
    """Initialize OpenAI client from environment variables."""
    api_base_url = os.getenv('API_BASE_URL', 'https://api.openai.com/v1')
    model_name = os.getenv('MODEL_NAME', 'gpt-4')
    
    # HF_TOKEN without default (required for Hugging Face)
    hf_token = os.getenv('HF_TOKEN')
    api_key = hf_token or os.getenv('API_KEY')
    
    if not api_key:
        raise ValueError("Neither HF_TOKEN nor API_KEY environment variables are set")
    
    # Initialize OpenAI client with base_url and api_key
    client = OpenAI(
        api_key=api_key,
        base_url=api_base_url
    )
    
    return client, model_name

def create_adaptive_prompt(observation: Dict[str, Any], step_type: str, agent_state: AgentState) -> str:
    """Create adaptive prompt based on agent learning state."""
    step_names = {
        'detect': 'Issue Detection',
        'risk': 'Risk Assessment', 
        'fix': 'Fix Suggestion'
    }
    
    step_instructions = {
        'detect': 'list all issues you find in the code',
        'risk': 'output the risk level (low/medium/high)',
        'fix': 'provide corrected code that fixes the identified issues'
    }
    
    # Adaptive prompt based on learning
    learning_context = ""
    if agent_state.past_failures:
        learning_context = f"\nLEARNING FROM PAST MISTAKES:\n{chr(10).join(f'- Avoid: {failure}' for failure in agent_state.past_failures[-3:])}"
    
    if agent_state.prompt_strategy == "detailed":
        instruction_detail = " Be thorough and provide detailed reasoning."
    elif agent_state.prompt_strategy == "focused":
        instruction_detail = " Focus on most critical issues first."
    else:
        instruction_detail = ""
    
    prompt = f"""You are a professional code reviewer learning from feedback.

Code:
{observation.code}

Intent:
{observation.intent}

Constraints:
{chr(10).join(f"- {constraint}" for constraint in observation.constraints)}

Step: {step_names[step_type]}

Instructions:
* {step_instructions[step_type]}{instruction_detail}
* Provide confidence score (0.0-1.0)
* Include your reasoning{learning_context}

Respond ONLY in JSON format:
{{
  "issues_detected": [...],
  "risk_level": "...",
  "fix": "...",
  "confidence": 0.0,
  "reasoning": "...",
  "reward_breakdown": {{
    "detect": 0.0,
    "risk": 0.0,
    "fix": 0.0
  }}
}}"""
    
    return prompt

def parse_llm_response(response_text: str) -> AgentResponse:
    """Parse LLM response into structured AgentResponse with fallback."""
    try:
        # Clean up response text
        response_text = response_text.strip()
        if response_text.startswith('```json'):
            response_text = response_text[7:]
        if response_text.endswith('```'):
            response_text = response_text[:-3]
        response_text = response_text.strip()
        
        parsed = json.loads(response_text)
        
        # Ensure all required fields exist with defaults
        return AgentResponse(
            issues_detected=parsed.get('issues_detected', []),
            risk_level=parsed.get('risk_level', 'low'),
            fix=parsed.get('fix', ''),
            confidence=max(0.0, min(1.0, parsed.get('confidence', 0.5))),
            reasoning=parsed.get('reasoning', ''),
            reward_breakdown=parsed.get('reward_breakdown', {})
        )
    except (json.JSONDecodeError, AttributeError, ValueError) as e:
        logger.warning(f"Failed to parse LLM response: {e}")
        logger.warning(f"Raw response: {response_text[:200]}...")
        
        # Return fallback response
        return AgentResponse(
            issues_detected=[],
            risk_level='low',
            fix='',
            confidence=0.1,
            reasoning='Failed to parse response',
            reward_breakdown={}
        )

def create_action_from_response(step_type: str, response: AgentResponse) -> Action:
    """Create Action model from structured AgentResponse."""
    return Action(
        action_type=step_type,
        detected_issues=response.issues_detected if step_type == 'detect' else None,
        risk_level=response.risk_level if step_type == 'risk' else None,
        suggested_fix=response.fix if step_type == 'fix' else None
    )

def run_llm_inference(client: OpenAI, model_name: str, prompt: str) -> str:
    """Run LLM inference with given prompt."""
    try:
        response = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": "You are a professional code reviewer. Always respond in valid JSON format with confidence and reasoning."},
                {"role": "user", "content": prompt}
            ],
            temperature=0,  # Deterministic
            max_tokens=500  # Increased for structured output
        )
        return response.choices[0].message.content
    except Exception as e:
        logger.error(f"LLM inference failed: {e}")
        raise

def run_baseline(env: CodeGuardianEnv) -> float:
    """Run a simple non-learning baseline agent for comparison."""
    print("STEP: Running baseline agent")
    logger.info("[BASELINE START] Running baseline agent...")
    
    obs = env.reset()
    total_reward = 0.0
    step_count = 0
    
    while not env.done and step_count < 50:
        current_step = 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 = env.step(action)
        total_reward += reward.value
        step_count += 1
        obs = next_obs
        
        logger.info(f"[BASELINE STEP] {step_count}: {step_type} -> reward={reward.value:+.3f}")
    
    print(f"STEP: Baseline complete - total_reward={total_reward:.3f} steps={step_count}")
    logger.info(f"[BASELINE END] Total reward: {total_reward:.3f}, Steps: {step_count}")
    return total_reward

def update_agent_state(agent_state: AgentState, episode_reward: float, feedback: str) -> AgentState:
    """Update agent state based on episode performance for policy improvement."""
    agent_state.episode += 1
    agent_state.total_reward = episode_reward
    agent_state.improvement_history.append(episode_reward)
    
    # Update best score
    if episode_reward > agent_state.best_score:
        agent_state.best_score = episode_reward
        logger.info(f"New best score achieved: {episode_reward:.3f}")
    
    # STRONGER POLICY UPDATE LOGIC
    # Treat poor performance as failure
    if episode_reward < agent_state.best_score * 0.8:
        agent_state.past_failures.append(f"Episode {agent_state.episode}: reward={episode_reward:.3f}")
        # Keep only recent failures
        agent_state.past_failures = agent_state.past_failures[-5:]
        logger.info(f"Failure recorded - Episode reward: {episode_reward:.3f} vs Best: {agent_state.best_score:.3f}")
    
    # IMPROVED STRATEGY SWITCHING
    if episode_reward < agent_state.best_score * 0.7:  # Poor performance threshold
        old_strategy = agent_state.prompt_strategy
        
        if agent_state.prompt_strategy == "default":
            agent_state.prompt_strategy = "detailed"
            logger.info("Switching strategy: default → detailed")
        elif agent_state.prompt_strategy == "detailed":
            agent_state.prompt_strategy = "focused"
            logger.info("Switching strategy: detailed → focused")
        else:
            # Reset to default if focused still fails
            agent_state.prompt_strategy = "default"
            logger.info("Switching strategy: focused → default (reset)")
        
        if old_strategy != agent_state.prompt_strategy:
            logger.info(f"Strategy change triggered by low performance: {episode_reward:.3f} < {agent_state.best_score * 0.7:.3f}")
    
    return agent_state

def run_episode(env: CodeGuardianEnv, client: OpenAI, model_name: str, agent_state: AgentState) -> tuple[float, List[Dict]]:
    """Run a single episode and return total reward with episode history."""
    obs = env.reset()
    episode_reward = 0.0
    episode_history = []
    step_count = 0
    
    print(f"START: Episode {agent_state.episode + 1} - model={model_name} task_id={obs.task_id} difficulty={obs.difficulty}")
    logger.info(f"[EPISODE START] Episode {agent_state.episode + 1} - model={model_name} task_id={obs.task_id} difficulty={obs.difficulty}")
    
    while not env.done:
        current_step = env.state()['current_step']
        step_types = ['detect', 'risk', 'fix']
        step_type = step_types[current_step]
        
        step_count += 1
        print(f"STEP: {step_count} - {step_type} - task_id={obs.task_id}")
        logger.info(f"[STEP] step_number={step_count} step_type={step_type} task_id={obs.task_id}")
        
        # Create adaptive prompt
        prompt = create_adaptive_prompt(obs, step_type, agent_state)
        
        # Run LLM inference
        try:
            response_text = run_llm_inference(client, model_name, prompt)
            logger.info(f"LLM response: {response_text[:100]}...")
        except Exception as e:
            logger.error(f"LLM inference failed: {e}")
            # Use fallback response
            agent_response = AgentResponse(
                issues_detected=[],
                risk_level='low',
                fix='',
                confidence=0.1,
                reasoning='Inference failed',
                reward_breakdown={}
            )
        else:
            # Parse structured response
            agent_response = parse_llm_response(response_text)
        
        # Create action from response
        action = create_action_from_response(step_type, agent_response)
        
        # Execute step
        next_obs, reward, done = env.step(action)
        
        # CONFIDENCE-BASED REWARD ADJUSTMENT
        confidence_penalty = (1 - agent_response.confidence) * 0.1
        adjusted_reward = reward.value - confidence_penalty
        episode_reward += adjusted_reward
        
        logger.info(f"confidence_penalty={confidence_penalty:.3f} adjusted_reward={adjusted_reward:+.3f} confidence={agent_response.confidence:.2f}")
        
        # Log step results
        step_info = {
            "step": step_count,
            "type": step_type,
            "action": action.dict(),
            "response": agent_response.dict(),
            "reward": reward.dict(),
            "adjusted_reward": adjusted_reward,
            "confidence_penalty": confidence_penalty,
            "total_reward": episode_reward
        }
        episode_history.append(step_info)
        
        logger.info(f"action_type={action.action_type} reward={reward.value:+.3f} adjusted_reward={adjusted_reward:+.3f} feedback={reward.feedback} total_reward={episode_reward:+.3f} current_step={env.state()['current_step']} current_task={env.state()['current_task']}")
        
        # Update observation for next step
        obs = next_obs
        
        # Safety check
        if step_count > 50:  # Prevent infinite loops
            logger.warning("Maximum steps reached, stopping episode")
            break
    
    # Final episode log
    print(f"END: Episode {agent_state.episode + 1} - final_score={episode_reward:.3f} total_steps={step_count} strategy={agent_state.prompt_strategy}")
    logger.info(f"[EPISODE END] Episode {agent_state.episode + 1} - final_score={episode_reward:.3f} total_steps={step_count} strategy={agent_state.prompt_strategy}")
    
    return episode_reward, episode_history

def train_agent(num_episodes: int = 10) -> AgentState:
    """Train agent over multiple episodes with policy improvement."""
    try:
        # Setup
        client, model_name = setup_openai_client()
        agent_state = AgentState()
        all_episodes = []
        
        print(f"START: RL training for {num_episodes} episodes with model: {model_name}")
        logger.info(f"[TRAINING START] Starting RL training for {num_episodes} episodes with model: {model_name}")
        
        # RUN BASELINE FIRST
        logger.info("\n{'='*60}")
        logger.info("BASELINE EVALUATION")
        logger.info(f"{'='*60}")
        baseline_env = CodeGuardianEnv()
        baseline_score = run_baseline(baseline_env)
        logger.info(f"Baseline score: {baseline_score:.3f}")
        
        # TRAINING LOOP
        logger.info(f"\n{'='*60}")
        logger.info("REINFORCEMENT LEARNING")
        logger.info(f"{'='*60}")
        
        for episode in range(num_episodes):
            logger.info(f"\n{'='*50}")
            logger.info(f"EPISODE {episode + 1}/{num_episodes}")
            logger.info(f"{'='*50}")
            
            env = CodeGuardianEnv()
            episode_reward, episode_history = run_episode(env, client, model_name, agent_state)
            all_episodes.append({
                "episode": episode + 1,
                "reward": episode_reward,
                "history": episode_history
            })
            
            # Update agent state for policy improvement
            agent_state = update_agent_state(agent_state, episode_reward, f"Episode {episode + 1} completed")
            
            logger.info(f"Episode {episode + 1} completed with reward: {episode_reward:.3f}")
            logger.info(f"Best score so far: {agent_state.best_score:.3f}")
            logger.info(f"Current strategy: {agent_state.prompt_strategy}")
            
            # EARLY STOPPING
            if episode > 3 and agent_state.best_score > 2.5:
                logger.info("Early stopping triggered: high performance achieved")
                logger.info(f"Stopping at episode {episode + 1} with best score {agent_state.best_score:.3f}")
                break
        
        # PERFORMANCE COMPARISON
        print("STEP: Training summary analysis")
        logger.info(f"\n{'='*60}")
        logger.info(f"[TRAINING SUMMARY]")
        logger.info(f"{'='*60}")
        
        # Calculate improvement over baseline
        improvement = ((agent_state.best_score - baseline_score) / (abs(baseline_score) + 1e-6)) * 100
        
        # Log key metrics
        logger.info(f"RL Agent Best Score: {agent_state.best_score:.3f}")
        logger.info(f"Baseline Score: {baseline_score:.3f}")
        logger.info(f"Performance Improvement: {improvement:.2f}%")
        
        # Additional metrics
        if all_episodes:
            avg_reward = sum(ep["reward"] for ep in all_episodes) / len(all_episodes)
            logger.info(f"Average Episode Reward: {avg_reward:.3f}")
            logger.info(f"Episodes Completed: {len(all_episodes)}")
            logger.info(f"Final Strategy: {agent_state.prompt_strategy}")
            
            # Performance classification
            if improvement > 50:
                print("STEP: EXCELLENT - significant improvement over baseline")
            elif improvement > 20:
                print("STEP: GOOD - measurable improvement achieved")
            elif improvement > 0:
                print("STEP: POSITIVE - some improvement over baseline")
            else:
                print("STEP: NO IMPROVEMENT - RL agent not better than baseline")
        
        return agent_state
        
    except Exception as e:
        logger.error(f"Training failed: {e}")
        raise

def main():
    """Main inference loop with RL training."""
    try:
        print("START")
        
        # Run RL training
        agent_state = train_agent(num_episodes=10)
        return 0
        
    except Exception as e:
        logger.error(f"Inference failed: {e}")
        return 1
    finally:
        print("END")

if __name__ == "__main__":
    exit(main())