import os
import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import logging
import csv
from datetime import datetime
from policy import RouterPolicy


# Initialize module logger
logger = logging.getLogger("ppo_trainer")
if not logger.handlers:
    logger.setLevel(logging.INFO)
    fmt = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
    # Console handler
    ch = logging.StreamHandler()
    ch.setFormatter(fmt)
    logger.addHandler(ch)
    # File handler (append) - place logs directory at repository root `logs/`
    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
    logs_dir = os.path.join(repo_root, 'logs')
    try:
        os.makedirs(logs_dir, exist_ok=True)
    except Exception:
        logs_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))

    fh_path = os.path.join(logs_dir, 'ppo_trainer.log')
    fh = logging.FileHandler(fh_path)
    fh.setFormatter(fmt)
    logger.addHandler(fh)


def _truncate(s: str, n: int = 500) -> str:
    if s is None:
        return ""
    s = str(s)
    return s if len(s) <= n else s[:n-1] + "\u2026"


def v1_ppo_update(model, optimizer, states, actions, old_log_probs, returns, advantages, clip_ratio=0.2):
    """
    Performs a single PPO update step.
    
    states: Tensor of query embeddings
    actions: Tensor of actions taken
    returns: Tensor of discounted cumulative rewards
    advantages: Tensor of advantages (Returns - State Values)
    """
    model.train()
    device = next(model.parameters()).device

    # ensure tensors live on the same device as the model
    states = states.to(device)
    actions = actions.to(device)
    old_log_probs = old_log_probs.to(device)
    returns = returns.to(device)
    advantages = advantages.to(device)

    #1. Get current policy probabilities and values
    logits, current_values = model(states)
    probs = F.softmax(logits, dim=-1)
    dist = torch.distributions.Categorical(probs)
    
    # calculate new log probabilities and entropy for the taken actions
    new_log_probs = dist.log_prob(actions)
    entropy = dist.entropy().mean()
    
    #2. calculate the ratio: pi_theta(a|s) / pi_theta_old(a|s)
    ratios = torch.exp(new_log_probs - old_log_probs)
    
    #3. calculate surrogate losses
    surr1 = ratios * advantages
    surr2 = torch.clamp(ratios, 1.0 - clip_ratio, 1.0 + clip_ratio) * advantages
    
    # calculate actor and critic losses
    # PPO loss is the negative of the minimum surrogate loss plus value function loss
    actor_loss = -torch.min(surr1, surr2).mean()  # PPO objective (maximize the minimum of the surrogate losses)
    critic_loss = F.mse_loss(current_values.squeeze(), returns.view_as(current_values.squeeze()))  # Value function loss
    
    #5. total Loss calculation (with entropy bonus to encourage exploration)
    entropy_coef = 0.08 #increased after v1 experiment #0.01 # common weight for the entropy bonus
    value_coef = 0.5 # common weight for the value loss
    total_loss = actor_loss + (value_coef * critic_loss) - (entropy_coef * entropy)  # Combine losses with appropriate weights
    
    #6. backpropagation and optimization step
    optimizer.zero_grad()
    total_loss.backward()
    
    # Gradient clipping (optional but common in PPO to stabilize training)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
    optimizer.step()
    
    return total_loss.item()
    
def ppo_update(model, optimizer, states, actions, old_log_probs, returns, advantages, current_epoch, total_epochs, clip_ratio=0.2):
    model.train()
    
    # 1. Forward pass
    logits, current_values = model(states)
    
    # Optional: Logit Temperature Scaling (helps soften the SFT spike)
    temperature = 1.5 
    probs = F.softmax(logits / temperature, dim=-1)
    dist = torch.distributions.Categorical(probs)
    
    new_log_probs = dist.log_prob(actions)
    entropy = dist.entropy().mean()
    
    # 2. Calculate PPO Ratio
    ratios = torch.exp(new_log_probs - old_log_probs)
    
    # 3. Calculate Surrogate Losses
    surr1 = ratios * advantages
    surr2 = torch.clamp(ratios, 1.0 - clip_ratio, 1.0 + clip_ratio) * advantages
    
    actor_loss = -torch.min(surr1, surr2).mean()
    critic_loss = F.mse_loss(current_values.squeeze(), returns)
    
    # ---------------------------------------------------------
    # NOVELTY FOR THE PAPER: Dynamic Entropy Annealing
    # Start high (0.2) to shatter SFT confidence, decay to low (0.01) to exploit
    # ---------------------------------------------------------
    initial_entropy = 0.2
    final_entropy = 0.01
    decay_rate = current_epoch / max(1, total_epochs - 1)
    dynamic_entropy_coef = initial_entropy - decay_rate * (initial_entropy - final_entropy)
    
    value_coef = 0.5
    total_loss = actor_loss + (value_coef * critic_loss) - (dynamic_entropy_coef * entropy)
    
    optimizer.zero_grad()
    total_loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
    optimizer.step()
    
    return total_loss.item(), dynamic_entropy_coef
    
def simulate_rag_pipeline(query, action_idx):
    """
    Dummy function mapping to your actual RAG modules.
    0: Unstructured, 1: Website, 2: Structured (SQL)
    """
    # Integrate with backend SimpleAgent to run the actual tools.
    # This function will synchronously run async agent methods using asyncio.run
    import sys
    import os
    import asyncio
    import requests

    # Ensure repo root is on sys.path so we can import backend package
    # repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
    # if repo_root not in sys.path:
    #     sys.path.insert(0, repo_root)

    import sys
    import os

    project_root = os.path.abspath(os.path.join(os.getcwd(), ".."))
    backend_path = os.path.join(project_root, "backend")

    sys.path.append(project_root)
    sys.path.append(backend_path)

    try:
        from backend.app.core.agent import SimpleAgent
    except Exception:
        # Fallback import path for different execution contexts
        try:
            from app.core.agent import SimpleAgent
        except Exception as e:
            return "", f"Failed importing SimpleAgent: {e}"
    def detect_language(text):
        #check if any character in the text is in the Korean Unicode range or chinese Unicode range or english Unicode range
        if any('\uac00' <= char <= '\ud7a3' for char in text):
            return "ko"
        elif any('\u4e00' <= char <= '\u9fff' for char in text):
            return "zh"
        else:
            return "en"
    
    async def old_call_tool(q, act_idx):
        agent = SimpleAgent()
        language = detect_language(q)
        if act_idx == 0:
            res = await agent._execute_guidelines(q, language)
        elif act_idx == 1:
            res = await agent._execute_web_search(q, language)
        else:
            res = await agent._execute_cbnu_database(q, language)

        # Build a short context from tool results
        context = ""
        try:
            results = res.data.get("results", []) if res and res.data else []
            snippets = []
            for item in results[:3]:
                snippet = item.get("snippet") or item.get("content") or item.get(f"answer_{language}") or str(item)
                snippets.append(snippet)
            context = "\n".join(snippets)
        except Exception:
            context = ""

        # Use agent's generator to synthesize an answer from the single tool result
        try:
            answer = await agent._generate_response(q, [res], language)
        except Exception:
            answer = ""

        return context, answer

    async def _call_tool(q, act_idx):
        language = detect_language(q)
        print(f"Detected language: {language}")
        
        # 0: guidelines, 1: web search, 2: database
        if act_idx == 0:
            search_sources = ["guidelines"]
        elif act_idx == 1:
            search_sources = ["web-search"]
        else:
            search_sources = ["course-registration", "department-info"]
            
        url = 'http://localhost:7000/api/chat'
        headers = {
            'accept': 'application/json',
            'Content-Type': 'application/json'
        }
        payload = {
            "message": q,
            "language": language,
            "search_sources": search_sources,
            "session_id": "ppo_train_session"
        }

        try:
            response = requests.post(url, headers=headers, json=payload)
            response.raise_for_status()
            data = response.json()
            
            # Extract context and answer from the API response
            answer = data.get("message", "")
            
            # Extract content securely from the sources
            context_parts = []
            sources = data.get("sources", [])
            for src in sources:
                if isinstance(src, dict) and src.get("success"):
                    src_data = src.get("data", {})
                    results = src_data.get("results", [])
                    for res in results:
                        if "content" in res:
                            context_parts.append(res["content"])
                        elif "snippet" in res:
                            context_parts.append(res["snippet"])
                            
            context = "\n".join(context_parts)
                
            print(f"Tool Result fetched successfully via API.")
            return context, answer

        except Exception as e:
            print(f"API call failed: {e}")
            return "", ""
    try:
        ctx, ans = asyncio.run(_call_tool(query, int(action_idx)))
        return ctx, ans
    except Exception as e:
        return "", f"simulate_rag_pipeline failed: {e}"

def llm_judge_reward(query, context, answer, action_idx):
    import os
    import requests
    import re
    
    try:
        # 1. FIXED: Correct API version (v1beta) and method (:generateContent)
        # Note: If gemini-2.5-flash still throws 404, change to gemini-2.0-flash or gemini-1.5-flash
        model_name = "gemini-2.5-flash" 
        gemini_url = f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent"
        action_taken = f"Action {action_idx}"
        gemini_key = "AIzaSyBoErDeUH1kRBbusI6qkMwISw5euyuXwZo"
        gemini_token = os.getenv("GEMINI_API_TOKEN") or os.getenv("GCP_ACCESS_TOKEN")
        
        Old_prompt = f"""You are an expert automated RAG evaluator. 
        Your task is to grade the retrieved Context and the generated Answer based on the user's Query.

        ### GRADING RUBRIC:
        1. Context Relevance (R_rel): Does the Context contain the necessary information to answer the Query?
        - 0.0: Completely irrelevant or useless.
        - 0.5: Contains partial information, but not enough for a complete answer.
        - 1.0: Contains all necessary information to fully answer the query.

        2. Answer Correctness (R_corr): Is the Answer fully supported by the Context? (IGNORE your internal knowledge. Grade strictly on the provided context).
        - 0.0: Hallucinates information not in the context, or completely fails to answer the query.
        - 0.5: Partially correct, but misses key details from the context or includes slight hallucinations.
        - 1.0: Perfectly accurate, directly answers the query, and is 100% faithful to the context.

        ### INPUTS:
        Query: {query}
        Context: {context}
        Answer: {answer}

        ### OUTPUT FORMAT:
        You must return ONLY a valid JSON object. Do not use markdown blocks (```json).
        Output this exact structure:
        {{
            "relevance_reasoning": "Brief 1-sentence explanation of why the context is/isn't relevant.",
            "R_rel": <float between 0.0 and 1.0>,
            "correctness_reasoning": "Brief 1-sentence explanation of answer correctness and faithfulness.",
            "R_corr": <float between 0.0 and 1.0>
        }}
        """

        prompt = f"""You are an expert automated RAG evaluator. 
            Your task is to grade the retrieved Context and the generated Answer based on the user's Query and the Action taken by the routing agent.

            ### STRICT ROUTING RULES:
            The agent MUST select the correct action based on the query intent:
            - Action 0 (guidelines): Required for academic information, university rules, visa, scholarships, graduation requirements, health insurance, and part-time job info.
            - Action 1 (web_search): Required for university news, announcements, and external information.
            - Action 2 (cbnu_database): Required for department info (contacts, phone, email, website) and course registration (course details, schedules, professors, credits).

            ### GRADING RUBRIC:
            1. Context Relevance (R_rel): 
            - FATAL PENALTY: Did the agent choose an action that violates the STRICT ROUTING RULES above? If YES, you MUST score R_rel as 0.0 immediately, regardless of the context.
            - 0.0: Violates routing rules OR context is completely irrelevant/useless.
            - 0.5: Follows rules, but context contains only partial information.
            - 1.0: Follows rules AND context contains all necessary information to fully answer the query.

            2. Answer Correctness (R_corr): Is the Answer fully supported by the Context? (IGNORE your internal knowledge. Grade strictly on the provided context).
            - 0.0: Hallucinates information not in the context, or completely fails to answer the query.
            - 0.5: Partially correct, but misses key details from the context or includes slight hallucinations.
            - 1.0: Perfectly accurate, directly answers the query, and is 100% faithful to the context.

            ### INPUTS:
            Query: {query}
            Action Taken: {action_taken}
            Context: {context}
            Answer: {answer}

            ### OUTPUT FORMAT:
            You must return ONLY a valid JSON object. Do not use markdown blocks (```json).
            Output this exact structure:
            {{
                "relevance_reasoning": "Brief 1-sentence explanation of why the rule was followed/broken and if context is relevant.",
                "R_rel": <float between 0.0 and 1.0>,
                "correctness_reasoning": "Brief 1-sentence explanation of answer correctness and faithfulness.",
                "R_corr": <float between 0.0 and 1.0>
            }}
            """
        
        headers = {"Content-Type": "application/json"}
        params = {}
        
        if gemini_key:
            params = {"key": gemini_key}
        elif gemini_token:
            headers["Authorization"] = f"Bearer {gemini_token}"

        # 2. FIXED: Correct Gemini payload structure
        payload = {
            "contents": [{
                "parts": [{"text": prompt}]
            }],
            "generationConfig": {
                # "maxOutputTokens": 200,
                "temperature": 0.0, # Deterministic output for consistent rewards
                "responseMimeType": "application/json" # 3. ADDED: Enforces JSON output
            }
        }

        resp = requests.post(gemini_url, params=params, json=payload, headers=headers, timeout=15)
        
        if resp.status_code == 404:
            print(f"Gemini endpoint returned 404. Check if '{model_name}' is available to your API key.", flush=True)
            print(f"Tried URL: {gemini_url}", flush=True)
            resp.raise_for_status()
            
        resp.raise_for_status()

        # 4. FIXED: Parsing the correct Gemini response path
        data = resp.json()
        text = ""
        try:
            text = data["candidates"][0]["content"]["parts"][0]["text"]
        except (KeyError, IndexError):
            text = resp.text # Fallback to raw text if structure is unexpected

        # Extract floats using your existing regex
        m_rel = re.search(r'"R_rel"\s*:\s*([0-9]*\.?[0-9]+)', text)
        m_corr = re.search(r'"R_corr"\s*:\s*([0-9]*\.?[0-9]+)', text)
        
        if m_rel and m_corr:
            R_rel = float(m_rel.group(1))
            R_corr = float(m_corr.group(1))
        else:
            nums = re.findall(r'([0-9]*\.?[0-9]+)', text)
            if len(nums) >= 2:
                R_rel = float(nums[0])
                R_corr = float(nums[1])
            else:
                R_rel, R_corr = 0.8, 0.9
                
    except Exception as e:
        print("LLM judge failed, using fallback reward. Error details:", flush=True)
        print(e, flush=True)
        R_rel, R_corr = 0.7, 0.7

    # The Reward Formulation: R_total = R_rel * (alpha + beta * R_corr)
    alpha, beta = 0.7, 0.3
    print(f"LLM Judge for '{query[:50]}...' Scores => R_rel: {R_rel}, R_corr: {R_corr}", flush=True)
    reward = R_rel * (alpha + beta * R_corr)

    return float(reward)


def llm_judge_reward_openai(query, context, answer, action_idx):
    import os
    import requests
    import re
    
    try:
        model_name = "gpt-4o" 
        openai_url = "https://api.openai.com/v1/chat/completions"
        action_taken = f"Action {action_idx}"
        openai_key = "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"
        if not openai_key:
            print("OPENAI_API_KEY not found in environment.", flush=True)
            return 0.0
            
        old_prompt = f"""You are an expert automated RAG evaluator. 
        Your task is to grade the retrieved Context and the generated Answer based on the user's Query.

        ### GRADING RUBRIC:
        1. Context Relevance (R_rel): Does the Context contain the necessary information to answer the Query?
        - 0.0: Completely irrelevant or useless.
        - 0.5: Contains partial information, but not enough for a complete answer.
        - 1.0: Contains all necessary information to fully answer the query.

        2. Answer Correctness (R_corr): Is the Answer fully supported by the Context? (IGNORE your internal knowledge. Grade strictly on the provided context).
        - 0.0: Hallucinates information not in the context, or completely fails to answer the query.
        - 0.5: Partially correct, but misses key details from the context or includes slight hallucinations.
        - 1.0: Perfectly accurate, directly answers the query, and is 100% faithful to the context.

        ### INPUTS:
        Query: {query}
        Context: {context}
        Answer: {answer}

        ### OUTPUT FORMAT:
        You must return ONLY a valid JSON object. Do not use markdown blocks (```json).
        Output this exact structure:
        {{
            "relevance_reasoning": "Brief 1-sentence explanation of why the context is/isn't relevant.",
            "R_rel": <float between 0.0 and 1.0>,
            "correctness_reasoning": "Brief 1-sentence explanation of answer correctness and faithfulness.",
            "R_corr": <float between 0.0 and 1.0>
        }}
        """
        prompt = f"""You are an expert automated RAG evaluator. 
            Your task is to grade the retrieved Context and the generated Answer based on the user's Query and the Action taken by the routing agent.

            ### STRICT ROUTING RULES:
            The agent MUST select the correct action based on the query intent:
            - Action 0 (guidelines): Required for academic information, university rules, visa, scholarships, graduation requirements, health insurance, and part-time job info.
            - Action 1 (web_search): Required for university news, announcements, and external information.
            - Action 2 (cbnu_database): Required for department info (contacts, phone, email, website) and course registration (course details, schedules, professors, credits).

            ### GRADING RUBRIC:
            1. Context Relevance (R_rel): 
            - FATAL PENALTY: Did the agent choose an action that violates the STRICT ROUTING RULES above? If YES, you MUST score R_rel as 0.0 immediately, regardless of the context.
            - 0.0: Violates routing rules OR context is completely irrelevant/useless.
            - 0.5: Follows rules, but context contains only partial information.
            - 1.0: Follows rules AND context contains all necessary information to fully answer the query.

            2. Answer Correctness (R_corr): Is the Answer fully supported by the Context? (IGNORE your internal knowledge. Grade strictly on the provided context).
            - 0.0: Hallucinates information not in the context, or completely fails to answer the query.
            - 0.5: Partially correct, but misses key details from the context or includes slight hallucinations.
            - 1.0: Perfectly accurate, directly answers the query, and is 100% faithful to the context.

            ### INPUTS:
            Query: {query}
            Action Taken: {action_taken}
            Context: {context}
            Answer: {answer}

            ### OUTPUT FORMAT:
            You must return ONLY a valid JSON object. Do not use markdown blocks (```json).
            Output this exact structure:
            {{
                "relevance_reasoning": "Brief 1-sentence explanation of why the rule was followed/broken and if context is relevant.",
                "R_rel": <float between 0.0 and 1.0>,
                "correctness_reasoning": "Brief 1-sentence explanation of answer correctness and faithfulness.",
                "R_corr": <float between 0.0 and 1.0>
            }}
            """

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {openai_key}"
        }
        
        payload = {
            "model": model_name,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.0,
            "response_format": {"type": "json_object"}
        }

        resp = requests.post(openai_url, json=payload, headers=headers, timeout=15)
        resp.raise_for_status()

        data = resp.json()
        text = data["choices"][0]["message"]["content"]

        # Extract floats
        m_rel = re.search(r'"R_rel"\s*:\s*([0-9]*\.?[0-9]+)', text)
        m_corr = re.search(r'"R_corr"\s*:\s*([0-9]*\.?[0-9]+)', text)
        
        if m_rel and m_corr:
            R_rel = float(m_rel.group(1))
            R_corr = float(m_corr.group(1))
        else:
            nums = re.findall(r'([0-9]*\.?[0-9]+)', text)
            if len(nums) >= 2:
                R_rel = float(nums[0])
                R_corr = float(nums[1])
            else:
                R_rel, R_corr = 0.8, 0.9
                
    except Exception as e:
        print("OpenAI LLM judge failed, using fallback reward. Error details:", flush=True)
        print(e, flush=True)
        R_rel, R_corr = 0.7, 0.7

    alpha, beta = 0.7, 0.3
    print(f"OpenAI Judge for '{query[:50]}...' Scores => R_rel: {R_rel}, R_corr: {R_corr}", flush=True)
    reward = R_rel * (alpha + beta * R_corr)

    return float(reward)




def load_router_policy(checkpoint_path: str, embedding_dim: int | None = None, device: str = "cpu") -> RouterPolicy:
    """Load a RouterPolicy wrapper from checkpoint path."""
    return RouterPolicy.load(checkpoint_path, embedding_dim=embedding_dim, device=device)


def save_policy_checkpoint(policy: RouterPolicy, path: str, optimizer: torch.optim.Optimizer | None = None, meta: dict | None = None):
    return policy.save_checkpoint(path, optimizer=optimizer, meta=meta)

def v1_run_ppo_rollout(model, embedder, queries, optimizer, epochs=5, device: str = "cpu"):
    """
    Executes the environment loop and updates the model.
    """
    model.train()
    gamma = 0.99 # Discount factor
    
    # Prepare CSV debug log (repo-level `logs/ppo_debug.csv`)
    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
    csv_path = os.path.abspath(os.path.join(repo_root, 'logs', 'ppo_debug.csv'))
    os.makedirs(os.path.dirname(csv_path), exist_ok=True)
    write_header = not os.path.exists(csv_path)
    csv_file = open(csv_path, mode='a', newline='', encoding='utf-8')
    csv_writer = csv.writer(csv_file)
    if write_header:
        csv_writer.writerow([
            'timestamp', 'epoch', 'step', 'query', 'action', 'action_log_prob', 'value', 'reward', 'context', 'answer', 'epoch_avg_reward', 'policy_loss'
        ])

    for epoch in range(epochs):
        states, actions, log_probs, rewards, values = [], [], [], [], []
        
        # 1. Rollout Phase: Gather Experience
        for step_idx, query in enumerate(queries):
            # Create embedding (SentenceTransformer may use inference_mode internally).
            state_emb = embedder.encode(query, convert_to_tensor=True)

            # Clone the embedding to ensure it's a regular tensor (not an "inference" tensor)
            # and move it to the model/device.
            try:
                state_emb = state_emb.clone().to(device)
            except Exception:
                # fallback: try model parameter device
                state_emb = state_emb.clone().to(next(model.parameters()).device)

            # Agent selects a module
            action, log_prob, value = model.get_action(state_emb)
            
            # Execute RAG and get LLM Judge Score
            context, answer = simulate_rag_pipeline(query, action.item())
            # reward = llm_judge_reward(query, context, answer, action.item())
            reward = llm_judge_reward_openai(query, context, answer, action.item())

            # Log per-step details to logger and CSV for debugging
            try:
                log_prob_val = float(log_prob.item()) if hasattr(log_prob, 'item') else float(log_prob)
            except Exception:
                log_prob_val = None
            try:
                value_val = float(value.item()) if hasattr(value, 'item') else float(value)
            except Exception:
                value_val = None

            logger.info("Step %d epoch %d | action=%s log_prob=%s value=%s reward=%.4f", step_idx, epoch, int(action.item()), log_prob_val, value_val, float(reward))
            csv_writer.writerow([
                datetime.utcnow().isoformat(), epoch, step_idx,
                _truncate(query, 800), int(action.item()), log_prob_val, value_val, float(reward),
                _truncate(context, 1000), _truncate(answer, 1000), '', ''
            ])
            
            # Store experience
            states.append(state_emb)
            actions.append(action)
            log_probs.append(log_prob)
            rewards.append(reward)
            values.append(value)
            
        # 2. Compute Returns and Advantages
        # (Simplified for single-step routing: Return = Reward)
        returns = torch.tensor(rewards, dtype=torch.float32).to(states[0].device)
        values_tensor = torch.cat(values).squeeze()
        advantages = returns - values_tensor.detach()
        
        # Normalize advantages for training stability
        if advantages.numel() > 1:
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        else:
            advantages = advantages - advantages.mean()
        
        # 3. PPO Update Phase
        # states_tensor = torch.stack(states)
        # actions_tensor = torch.stack(actions)
        # log_probs_tensor = torch.stack(log_probs)
        states_tensor = torch.stack(states).detach()
        actions_tensor = torch.stack(actions).detach()
        log_probs_tensor = torch.stack(log_probs).detach() 
        
        # Ensure these are also detached (just to be safe)
        returns = returns.detach()
        advantages = advantages.detach()
        
        
        # loss = ppo_update(
        #     model, optimizer, states_tensor, actions_tensor, 
        #     log_probs_tensor, returns, advantages
        # )
        
        # NEW: Run multiple optimization passes over the same rollout data
        ppo_inner_epochs = 4
        for _ in range(ppo_inner_epochs):
            loss = ppo_update(
                model, optimizer, states_tensor, actions_tensor, 
                log_probs_tensor, returns, advantages
            )
        

        avg_reward = float(returns.mean())
        logger.info("PPO Epoch %d | Avg Reward: %.4f | Policy Loss: %.4f", epoch+1, avg_reward, loss)
        csv_writer.writerow([
            datetime.utcnow().isoformat(), epoch, '', '', '', '', '', '', '', '', avg_reward, loss
        ])
    
    try:
        csv_file.flush()
        csv_file.close()
    except Exception:
        pass
        
    print(f"PPO Epoch {epoch+1} | Avg Reward: {returns.mean():.4f} | Policy Loss: {loss:.4f}")
    
    
def run_ppo_rollout(model, embedder, queries, optimizer, epochs=40, device: str = "cpu"):
    model.train()
    
    # Prepare CSV debug log (repo-level `logs/ppo_debug.csv`)
    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
    csv_path = os.path.abspath(os.path.join(repo_root, 'logs', 'ppo_debug.csv'))
    os.makedirs(os.path.dirname(csv_path), exist_ok=True)
    write_header = not os.path.exists(csv_path)
    csv_file = open(csv_path, mode='a', newline='', encoding='utf-8')
    csv_writer = csv.writer(csv_file)
    if write_header:
        csv_writer.writerow([
            'timestamp', 'epoch', 'step', 'query', 'action', 'action_log_prob', 'value', 'reward', 'context', 'answer', 'epoch_avg_reward', 'policy_loss'
        ])
    
    for epoch in range(epochs):
        states, actions, log_probs, rewards, values = [], [], [], [], []
        
        # 1. Rollout (Data Gathering)
        # for query in queries:
        for step_idx, query in enumerate(queries):
            with torch.no_grad():
                state_emb = embedder.encode(query, convert_to_tensor=True)
            
            # Clone and move to device
            try:
                state_emb = state_emb.clone().to(device)
            except Exception:
                state_emb = state_emb.clone().to(next(model.parameters()).device)
                
            action, log_prob, value = model.get_action(state_emb)
            context, answer = simulate_rag_pipeline(query, action.item())
            # reward = llm_judge_reward(query, context, answer, action.item())
            reward = llm_judge_reward_openai(query, context, answer, action.item())
            
            # Log per-step details to logger and CSV for debugging
            try:
                log_prob_val = float(log_prob.item()) if hasattr(log_prob, 'item') else float(log_prob)
            except Exception:
                log_prob_val = None
            try:
                value_val = float(value.item()) if hasattr(value, 'item') else float(value)
            except Exception:
                value_val = None
                
            logger.info("Step %d epoch %d | action=%s log_prob=%s value=%s reward=%.4f", step_idx, epoch, int(action.item()), log_prob_val, value_val, float(reward))
            csv_writer.writerow([
                datetime.utcnow().isoformat(), epoch, step_idx,
                _truncate(query, 800), int(action.item()), log_prob_val, value_val, float(reward),
                _truncate(context, 1000), _truncate(answer, 1000), '', ''
            ])
            
            states.append(state_emb)
            actions.append(action)
            log_probs.append(log_prob)
            rewards.append(reward)
            values.append(value)
            
            
        returns = torch.tensor(rewards, dtype=torch.float32).to(states[0].device)
        values_tensor = torch.cat(values).squeeze()
        advantages = returns - values_tensor.detach()
        
        # ---------------------------------------------------------
        # CRITICAL FIX: Strict Advantage Normalization
        # Prevents positive-only rewards from causing gradient collapse
        # ---------------------------------------------------------
        # advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        if advantages.numel() > 1:
            advantages = (advantages - advantages.mean()) / (advantages.std(unbiased=False) + 1e-8)
        else:
            advantages = advantages - advantages.mean()
        
        states_tensor = torch.stack(states).detach()
        actions_tensor = torch.stack(actions).detach()
        log_probs_tensor = torch.stack(log_probs).detach()
        returns = returns.detach()
        advantages = advantages.detach()
        
        # 2. PPO Inner Optimization Loop
        ppo_inner_epochs = 4 
        epoch_loss = 0
        
        for _ in range(ppo_inner_epochs):
            loss, current_ent_coef = ppo_update(
                model, optimizer, states_tensor, actions_tensor, 
                log_probs_tensor, returns, advantages,
                current_epoch=epoch, total_epochs=epochs
            )
            epoch_loss += loss
        avg_reward = float(returns.mean())
        logger.info("PPO Epoch %d | Avg Reward: %.4f | Policy Loss: %.4f", epoch+1, avg_reward, loss)
        csv_writer.writerow([
            datetime.utcnow().isoformat(), epoch, '', '', '', '', '', '', '', '', avg_reward, loss
        ])
            
        print(f"Epoch {epoch+1}/{epochs} | Avg Reward: {returns.mean():.4f} | Ent Coef: {current_ent_coef:.4f} | Loss: {epoch_loss/ppo_inner_epochs:.4f}")