#!/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__)

# Module-level environment variables
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")

DEFAULT_API_BASE_URL = "https://router.huggingface.co/v1"
DEFAULT_MODEL_NAME = "meta-llama/Meta-Llama-3-8B-Instruct"
HF_FALLBACK_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"


def _hf_auth_tokens() -> Optional[str]:
    return os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")


def _openai_auth_tokens() -> Optional[str]:
    return os.getenv("OPENAI_API_KEY") or os.getenv("API_KEY")


def _looks_like_huggingface_endpoint(url: str) -> bool:
    u = (url or "").lower()
    return any(x in u for x in ("huggingface", "hf.co", "router.huggingface"))


def _looks_like_openai_endpoint(url: str) -> bool:
    return "openai.com" in (url or "").lower()


def _looks_like_hf_hub_model(model: str) -> bool:
    """Hugging Face router / Inference typically uses org/model ids."""
    return "/" in (model or "").strip()


def _looks_like_openai_model(model: str) -> bool:
    m = (model or "").strip().lower()
    if not m or "/" in m:
        return False
    prefixes = (
        "gpt-",
        "o1",
        "o3",
        "o4",
        "ft:",
        "davinci",
        "text-embedding",
        "chatgpt-",
    )
    return any(m.startswith(p) for p in prefixes)


def _resolve_api_key(api_base_url: str, model_name: str) -> str:
    """
    Pick API key for an OpenAI-compatible client: Hugging Face token and/or OpenAI key,
    ordered from URL + model hints, with 'or' fallbacks so either stack can work.
    """
    hf_key = _hf_auth_tokens()
    oai_key = _openai_auth_tokens()

    url_hf = _looks_like_huggingface_endpoint(api_base_url)
    url_oai = _looks_like_openai_endpoint(api_base_url)
    model_hf = _looks_like_hf_hub_model(model_name)
    model_oai = _looks_like_openai_model(model_name)

    prefer_hf = url_hf or model_hf
    prefer_oai = url_oai or model_oai

    key: Optional[str] = None
    if prefer_hf and not prefer_oai:
        key = hf_key or oai_key
    elif prefer_oai and not prefer_hf:
        key = oai_key or hf_key
    elif prefer_hf and prefer_oai:
        if url_hf and not url_oai:
            key = hf_key or oai_key
        elif url_oai and not url_hf:
            key = oai_key or hf_key
        else:
            key = hf_key or oai_key
    else:
        key = hf_key or oai_key

    if not key:
        raise ValueError(
            "Set authentication for your provider: "
            "(HF_TOKEN or HUGGINGFACEHUB_API_TOKEN) for Hugging Face, "
            "or (OPENAI_API_KEY or API_KEY) for OpenAI"
        )
    return key


def _fallback_llm_json_string(reason: str = "LLM unavailable") -> str:
    return json.dumps(
        {
            "issues_detected": [],
            "risk_level": "low",
            "fix": "",
            "confidence": 0.1,
            "reasoning": reason,
            "reward_breakdown": {"detect": 0.0, "risk": 0.0, "fix": 0.0},
        }
    )


def setup_openai_client() -> tuple[OpenAI, str]:
    """Initialize OpenAI client from environment variables."""

    print("STEP: Initializing OpenAI client")

    api_base_url = os.getenv("API_BASE_URL", DEFAULT_API_BASE_URL)
    model_name = os.getenv("MODEL_NAME", DEFAULT_MODEL_NAME)
    api_key = _resolve_api_key(api_base_url, model_name)

    hf_key = _hf_auth_tokens()
    oai_key = _openai_auth_tokens()
    if api_key == hf_key and hf_key:
        auth_label = "Hugging Face token (HF_TOKEN / HUGGINGFACEHUB_API_TOKEN)"
    elif api_key == oai_key and oai_key:
        auth_label = "OpenAI key (OPENAI_API_KEY / API_KEY)"
    else:
        auth_label = "resolved API key"

    print(f"DEBUG: API_BASE_URL={api_base_url}")
    print(f"DEBUG: MODEL_NAME={model_name}")
    print(f"DEBUG: auth={auth_label}")
    logger.info(
        "OpenAI-compatible client: API_BASE_URL=%s MODEL_NAME=%s auth=%s",
        api_base_url,
        model_name,
        auth_label,
    )

    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 _is_model_missing_error(exc: BaseException) -> bool:
    msg = str(exc).lower()
    if "model_not_found" in msg or "model not found" in msg:
        return True
    body = getattr(exc, "body", None)
    if isinstance(body, str) and "model_not_found" in body.lower():
        return True
    if isinstance(body, dict):
        err = body.get("error")
        if isinstance(err, dict) and err.get("code") == "model_not_found":
            return True
    return False


def run_llm_inference(client: OpenAI, model_name: str, prompt: str) -> str:
    """Run LLM inference; never raises — returns JSON text or a safe fallback."""
    models_to_try: List[str] = [model_name]
    if HF_FALLBACK_MODEL and HF_FALLBACK_MODEL not in models_to_try:
        models_to_try.append(HF_FALLBACK_MODEL)

    last_error: Optional[BaseException] = None
    for m in models_to_try:
        try:
            response = client.chat.completions.create(
                model=m,
                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,
                max_tokens=500,
            )
            content = response.choices[0].message.content
            if content is None or not str(content).strip():
                logger.error("LLM returned empty content for model=%s", m)
                last_error = ValueError("empty LLM content")
                continue
            if m != model_name:
                logger.info("LLM succeeded using fallback model=%s (primary=%s failed)", m, model_name)
            return str(content)
        except Exception as e:
            last_error = e
            logger.error(
                "LLM inference failed for model=%s: %s",
                m,
                e,
                exc_info=not _is_model_missing_error(e),
            )
            if _is_model_missing_error(e):
                logger.error("model_not_found for model=%s; will try next model if available", m)
            continue

    logger.error("All LLM models exhausted; last_error=%s", last_error)
    return _fallback_llm_json_string("LLM call failed — using fallback response")

def run_baseline(env: CodeGuardianEnv) -> float:
    """Run a simple non-learning baseline agent for comparison."""
    print("STEP: 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}")
    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"STEP: Episode {agent_state.episode + 1} started - 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_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 (never raises)
        response_text = run_llm_inference(client, model_name, prompt)
        logger.info("LLM response (truncated): %s...", response_text[:100])
        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"STEP: Episode {agent_state.episode + 1} complete - 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"STEP: RL training started - episodes={num_episodes} 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")
        
        LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
        if LOCAL_IMAGE_NAME:
            print(f"STEP: inference using Docker image {LOCAL_IMAGE_NAME}")
        else:
            print("STEP: Starting inference")
        
        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())