src/routes/code_routes.py

import io
import logging
import re
import json
from fastapi import APIRouter, Depends, HTTPException, Request
from typing import Dict, Optional, List, Tuple
from pydantic import BaseModel

from scripts.format_response import execute_code_from_markdown, format_code_block
from src.utils.logger import Logger
from src.routes.session_routes import get_session_id_dependency
from src.agents.agents import code_edit, code_fix
from src.db.schemas.models import CodeExecution
from src.db.init_db import get_session
import dspy
import textwrap
import os
from src.schemas.code_schema import CodeExecuteRequest, CodeEditRequest, CodeFixRequest, CodeCleanRequest, GetLatestCodeRequest
from src.utils.model_registry import MODEL_OBJECTS
import asyncio
import traceback

def clean_print_statements(code_block):
    """
    This function cleans up any `print()` statements that might contain unwanted `\n` characters.
    It ensures print statements are properly formatted without unnecessary newlines.
    """
    # This regex targets print statements, even if they have newlines inside
    return re.sub(r'print\((.*?)(\\n.*?)(.*?)\)', r'print(\1\3)', code_block, flags=re.DOTALL)


def remove_main_block(code):
    # Match the __main__ block
    pattern = r'(?m)^if\s+__name__\s*==\s*["\']__main__["\']\s*:\s*\n((?:\s+.*\n?)*)'
    
    match = re.search(pattern, code)
    if match:
        main_block = match.group(1)
        
        # Dedent the code block inside __main__
        dedented_block = textwrap.dedent(main_block)
        
        # Remove \n from any print statements in the block (also handling multiline print cases)
        dedented_block = clean_print_statements(dedented_block)
        # Replace the block in the code
        cleaned_code = re.sub(pattern, dedented_block, code)
        
        # Optional: Remove leading newlines if any
        cleaned_code = cleaned_code.strip()
        
        return cleaned_code
    return code

# Initialize router
router = APIRouter(
    prefix="/code",
    tags=["code"],
    responses={404: {"description": "Not found"}},
)

# Initialize logger
logger = Logger("code_routes", see_time=True, console_log=False)
try_logger = Logger("try_code_routes", see_time=True, console_log=False)


def score_code(args, code, datasets=None):
    """
    Simple code scorer that checks if code runs successfully.
    Handles both deep analysis (combined_code) and fix (fixed_code) scenarios.
    
    Args:
        args: Arguments (unused but required for dspy.Refine)
        code: Code object with combined_code, fixed_code, or string content
        datasets: Dictionary of datasets from session state (optional)
    
    Returns:
        int: Score (0=error, 1=success)
    """
    try:
        # Handle different attribute names based on context
        if hasattr(code, 'combined_code'):
            code_to_execute = code.combined_code  # Deep analysis
        elif hasattr(code, 'fixed_code'):
            code_to_execute = code.fixed_code     # Code fix
        else:
            code_to_execute = str(code)           # Fallback to string
        
        # Make datasets available if provided
        local_vars = {}
        if datasets:
            local_vars.update(datasets)
            
        # Execute the code to test if it works
        exec(code_to_execute, local_vars)
        return 1  # Success
    except Exception as e:
        return 0  # Error
   

# Remove the global refine_fixer declaration


def format_code(code: str) -> str:
    """
    Clean the code by organizing imports and ensuring code blocks are properly formatted.
    
    Args:
        code (str): The raw Python code as a string.
        
    Returns:
        str: The cleaned code.
    """
    # Move imports to top
    code = move_imports_to_top(code)
    
    # Split code into blocks if they exist (based on comments like '# agent_name code start')
    code_blocks = []
    current_block = []
    current_agent = None
    
    for line in code.splitlines():
        if re.search(r'#\s+\w+\s+code\s+start', line.lower()):
            if current_agent and current_block:
                code_blocks.append((current_agent, '\n'.join(current_block)))
                current_block = []
            current_agent = re.search(r'#\s+(\w+)\s+code\s+start', line.lower()).group(1)
            current_block.append(line)
        elif re.search(r'#\s+\w+\s+code\s+end', line.lower()):
            if current_block:
                current_block.append(line)
                code_blocks.append((current_agent, '\n'.join(current_block)))
                current_agent = None
                current_block = []
        else:
            current_block.append(line)
    
    # If there's remaining code not in a block
    if current_block:
        if current_agent:
            code_blocks.append((current_agent, '\n'.join(current_block)))
        else:
            code_blocks.append(('main', '\n'.join(current_block)))
    
    # If no blocks were identified, return the original cleaned code
    if not code_blocks:
        return code
    # Reconstruct the code with the identified blocks
    return '\n\n'.join([block[1] for block in code_blocks])

def extract_code_blocks(code: str) -> Dict[str, str]:
    """
    Extract code blocks from the code based on agent name comments.
    
    Args:
        code (str): The code containing multiple blocks
        
    Returns:
        Dict[str, str]: Dictionary mapping agent names to their code blocks
    """
    # Find code blocks with start and end markers
    block_pattern = r'(#\s+(\w+)\s+code\s+start[\s\S]*?#\s+\w+\s+code\s+end)'
    blocks_with_markers = re.findall(block_pattern, code, re.DOTALL)
    
    if not blocks_with_markers:
        # If no blocks found, treat the entire code as one block
        return {'main': code}
    
    result = {}
    for full_block, agent_name in blocks_with_markers:
        result[agent_name.lower()] = full_block.strip()
    
    return result

def identify_error_blocks(code: str, error_output: str) -> List[Tuple[str, str, str]]:
    """
    Identify code blocks that have errors during execution.
    
    Args:
        code (str): The full code containing multiple agent blocks
        error_output (str): The error output from execution
        
    Returns:
        List[Tuple[str, str, str]]: List of tuples containing (agent_name, block_code, error_message)
    """
    # Parse the error output to find which agents had errors
    faulty_blocks = []
    
    # Find error patterns like "=== ERROR IN AGENT_NAME ===" or "=== ERROR IN UNKNOWN_AGENT ==="
    error_matches = []
    for match in re.finditer(
        r'^===\s+ERROR\s+IN\s+([A-Za-z0-9_]+)\s+===\s*([\s\S]*?)(?=^===\s+[A-Z]+\s+IN\s+[A-Za-z0-9_]+\s+===|\Z)',
        error_output,
        re.MULTILINE
    ):
        error_matches.append((match.group(1), match.group(2)))
    
    if not error_matches:
        return []
    
    # Find all code blocks in the given code
    blocks = {}
    for agent_match in re.finditer(r'#\s+(\w+)\s+code\s+start([\s\S]*?)#\s+\w+\s+code\s+end', code, re.DOTALL):
        agent_name = agent_match.group(1).lower()
        full_block = agent_match.group(0)
        blocks[agent_name] = full_block
    
    # Match errors with their corresponding code blocks
    matched_blocks = set()
    for agent_name, error_message in error_matches:
        # Format from error output is AGENT_NAME_AGENT, we need to extract the base name
        # Remove '_AGENT' suffix if present and convert to lowercase
        normalized_name = agent_name.lower()
        if normalized_name.endswith('_agent'):
            normalized_name = normalized_name[:-6]  # Remove '_agent' suffix
        
        # Try direct match first
        if normalized_name in blocks:
            # Extract the relevant error information
            processed_error = extract_relevant_error_section(error_message)
            faulty_blocks.append((normalized_name, blocks[normalized_name], processed_error))
            matched_blocks.add(normalized_name)
        else:
            # Try fuzzy matching for agent names
            for block_name, block_code in blocks.items():
                if block_name not in matched_blocks and (normalized_name in block_name or block_name in normalized_name):
                    # Extract the relevant error information
                    processed_error = extract_relevant_error_section(error_message)
                    faulty_blocks.append((block_name, block_code, processed_error))
                    matched_blocks.add(block_name)
                    break
    
    # logger.log_message(f"Faulty blocks found: {len(faulty_blocks)}", level=logging.INFO)
    # logger.log_message(f"Faulty blocks: {faulty_blocks}", level=logging.INFO)
    return faulty_blocks

def extract_relevant_error_section(error_message: str) -> str:
    """
    Extract the most relevant parts of the error message to help with fixing.
    
    Args:
        error_message (str): The full error message
        
    Returns:
        str: The processed error message with the most relevant information
    """
    error_lines = error_message.strip().split('\n')
    
    # If "Problem at this location" is in the error, focus on that section
    if 'Problem at this location:' in error_message:
        problem_idx = -1
        for i, line in enumerate(error_lines):
            if 'Problem at this location:' in line:
                problem_idx = i
                break
        
        if problem_idx >= 0:
            # Include the "Problem at this location" section and a few lines after
            end_idx = min(problem_idx + 10, len(error_lines))
            problem_section = error_lines[problem_idx:end_idx]
            
            # Also include the error type from the end
            error_type_lines = []
            for line in reversed(error_lines):
                if line.startswith('TypeError:') or line.startswith('ValueError:') or line.startswith('AttributeError:'):
                    error_type_lines = [line]
                    break
            
            return '\n'.join(problem_section + error_type_lines)
    
    # If we couldn't find "Problem at this location", include first few and last few lines
    if len(error_lines) > 10:
        return '\n'.join(error_lines[:5] + error_lines[-7:])
    
    # If the error is short enough, return as is
    return error_message

async def fix_code_with_dspy(code: str, error: str, dataset_context: str = "", datasets: dict = None):
    """
    Fix code using DSPy Refine with datasets-aware reward function
    """
    try:
        # Wrap score_code to fix datasets argument
        reward_fn_with_datasets = lambda args, pred: score_code(args, pred, datasets=datasets)

        refine_fixer = dspy.Refine(
            module=dspy.Predict(code_fix),
            N=3,
            threshold=1.0,
            reward_fn=reward_fn_with_datasets,
            fail_count=3
        )
        
        # Check if we have valid API key
        anthropic_key = os.environ.get('ANTHROPIC_API_KEY')
        if not anthropic_key:
            raise ValueError("ANTHROPIC_API_KEY environment variable is not set")
        
        # Fix the entire code using refine
        try:
            # Create the LM instance that will be used
            thread_lm = MODEL_OBJECTS['claude-sonnet-4-6']
            
            # Define the blocking function to run in thread
            def run_refine_fixer():
                with dspy.context(lm=thread_lm):
                    return refine_fixer(
                        dataset_context=str(dataset_context) or "",
                        faulty_code=str(code) or "",
                        error=str(error) or "",
                    )
            
            # Use asyncio.to_thread for better async integration with timeout
            result = await asyncio.wait_for(
                asyncio.to_thread(run_refine_fixer), 
                timeout=60.0  # 60 second timeout
            )
            
            if not hasattr(result, 'fixed_code'):
                raise ValueError("DSPy Refine did not return a result with 'fixed_code' attribute")
            
            return result.fixed_code
            
        except Exception as e:
            logger.log_message(f"🔧 DETAILED ERROR in fix_code_with_dspy: {str(e)}", level=logging.ERROR)
            logger.log_message(f"�� ERROR TYPE: {type(e).__name__}", level=logging.ERROR)
            logger.log_message(f"�� ERROR TRACEBACK: {traceback.format_exc()}", level=logging.ERROR)
            
            # Instead of returning original code, raise the error so we can see what's wrong
            raise HTTPException(status_code=500, detail=f"Fix failed: {str(e)}")
            
    except Exception as e:
        logger.log_message(f"Error in fix_code_with_dspy: {str(e)}", level=logging.ERROR)
        raise RuntimeError(f"Fix code setup failed: {str(e)}") from e

def get_dataset_context(df):
    """
    Generate context information about the dataset
    
    Args:
        df: The pandas dataframe
         
    Returns:
        String with dataset information (columns, types, null values)
    """
    if df is None:
        return "No dataset is currently loaded."
    
    try:
        # Get basic dataframe info
        col_types = df.dtypes.to_dict()
        null_counts = df.isnull().sum().to_dict()
        
        # Format the context string
        context = "Dataset context:\n"
        context += f"- Shape: {df.shape[0]} rows, {df.shape[1]} columns\n"
        context += "- Columns and types:\n"
        
        for col, dtype in col_types.items():
            null_count = null_counts.get(col, 0)
            null_percent = (null_count / len(df)) * 100 if len(df) > 0 else 0
            context += f"  * {col} ({dtype}): {null_count} null values\n"
        
        # Add sample values for each column (first 2 non-null values)
        context += "- Sample values:\n"
        for col in df.columns:
            sample_values = df[col].dropna().head(2).tolist()
            # if float, round to 2 decimal places
            if df[col].dtype == "float64":
                sample_values = [round(v, 1) for v in sample_values]
            sample_str = ", ".join(str(v) for v in sample_values)
            context += f"  * {col}: {sample_str}\n"
        return context
    except Exception as e:
        return "Could not generate dataset context information."

def edit_code_with_dspy(original_code: str, user_prompt: str, dataset_context: str = ""):
    thread_lm = MODEL_OBJECTS['claude-sonnet-4-6']
    with dspy.context(lm=thread_lm):
        code_editor = dspy.Predict(code_edit)
        
        result = code_editor(
            dataset_context=dataset_context,
            original_code=original_code,
            user_prompt=user_prompt,
        )
        return result.edited_code

def move_imports_to_top(code: str) -> str:
    """
    Moves all import statements to the top of the Python code.

    Args:
        code (str): The raw Python code as a string.

    Returns:
        str: The cleaned code with import statements at the top.
    """
    # Extract import statements
    import_statements = re.findall(
        r'^\s*(import\s+[^\n]+|from\s+[^\n]+import\s+[^\n]+)', code, flags=re.MULTILINE
    )
    
    # Remove import statements from original code
    code_without_imports = re.sub(
        r'^\s*(import\s+[^\n]+|from\s+[^\n]+import\s+[^\n]+)\n?', '', code, flags=re.MULTILINE
    )
    
    # Deduplicate and sort imports
    sorted_imports = sorted(set(import_statements))
    
    # Combine cleaned imports and remaining code
    cleaned_code = '\n'.join(sorted_imports) + '\n\n' + code_without_imports.strip()
    
    return cleaned_code


@router.post("/execute")
async def execute_code(
    request_data: CodeExecuteRequest,
    request: Request,
    session_id: str = Depends(get_session_id_dependency)
):
    """
    Execute code provided in the request against the session's dataframe
    
    Args:
        request_data: Body containing code to execute
        request: FastAPI Request object
        session_id: Session identifier
        
    Returns:
        Dictionary containing execution output and any plot outputs
    """
    # Access app state via request
    app_state = request.app.state
    session_state = app_state.get_session_state(session_id)
    # logger.log_message(f"Session State: {session_state}", level=logging.INFO)
    
    if session_state["datasets"] is None:
        raise HTTPException(
            status_code=400,
            detail="No dataset is currently loaded. Please link a dataset before executing code."
        )
        
    try:
        code = request_data.code
        if not code:
            raise HTTPException(status_code=400, detail="No code provided")
            
        # Get the user_id and chat_id from session state if available
        user_id = session_state.get("user_id")
        chat_id = session_state.get("chat_id")
        message_id = request_data.message_id
        
        # If message_id was not provided in the request, try to get it from the session state
        if message_id is None:
            message_id = session_state.get("current_message_id")
        else:
            # Update the session state with the provided message_id
            session_state["current_message_id"] = message_id
        
        # Get model configuration
        model_config = session_state.get("model_config", {})
        model_provider = model_config.get("provider", "")
        model_name = model_config.get("model", "")
        model_temperature = model_config.get("temperature", 0.0)
        model_max_tokens = model_config.get("max_tokens", 0)
        
        # Get database session
        db = get_session()
        
        # Check if we have an existing execution record for this message
        existing_execution = None
        if message_id:
            try:
                existing_execution = db.query(CodeExecution).filter(
                    CodeExecution.message_id == message_id
                ).first()
                
            except Exception as query_error:
                logger.log_message(f"Error querying for existing execution: {str(query_error)}", level=logging.ERROR)
                # Continue without existing execution
        else:
            logger.log_message("No message_id provided in session state", level=logging.WARNING)
            
        # Execute the code with the dataframe from session state
        full_output = ""
        json_outputs = []
        matplotlib_outputs = []
        is_successful = True
        failed_agents = None
        error_messages = None
        
        try:
            full_output, json_outputs, matplotlib_outputs = execute_code_from_markdown(code, session_state["datasets"])
            
            # Even with "successful" execution, check for agent failures in the output
            failed_blocks = identify_error_blocks(code, full_output)
            
            if failed_blocks:
                # We have some failed agents even though no exception was thrown
                is_successful = False  # Mark as failed if any agent failed
                failed_agents = json.dumps([block[0] for block in failed_blocks])
                error_messages = json.dumps({
                    block[0]: block[2] for block in failed_blocks
                })
                logger.log_message(f"Partial execution failure. Failed agents: {failed_agents}", level=logging.WARNING)
            
        except Exception as exec_error:
            full_output = str(exec_error)
            json_outputs = []
            matplotlib_outputs = []
            is_successful = False
            
            # Identify which agents failed
            failed_blocks = identify_error_blocks(code, full_output)
            
            # Format the failed agents and error messages
            if failed_blocks:
                failed_agents = json.dumps([block[0] for block in failed_blocks])
                error_messages = json.dumps({
                    block[0]: block[2] for block in failed_blocks
                })
                logger.log_message(f"Execution threw exception. Failed agents: {failed_agents}", level=logging.ERROR)
            
            # Don't re-raise the error - we want to capture the error and send it back to the client
            # return error details in the response instead
        
        # Create or update the execution record regardless of success/failure
        try:
            if existing_execution:
                # Update existing record
                existing_execution.latest_code = code
                existing_execution.is_successful = is_successful
                existing_execution.output = full_output
                
                if not is_successful:
                    existing_execution.failed_agents = failed_agents
                    existing_execution.error_messages = error_messages
                    
                db.commit()
            else:
                # Create new record
                logger.log_message(f"Creating new CodeExecution record for message_id: {message_id}", level=logging.INFO)
                new_execution = CodeExecution(
                    message_id=message_id,
                    chat_id=chat_id,
                    user_id=user_id,
                    initial_code=code,
                    latest_code=code,
                    is_successful=is_successful,
                    output=full_output,
                    model_provider=model_provider,
                    model_name=model_name,
                    model_temperature=model_temperature,
                    model_max_tokens=model_max_tokens,
                    failed_agents=failed_agents,
                    error_messages=error_messages
                )
                db.add(new_execution)
                db.commit()
                logger.log_message(f"Successfully created CodeExecution record with ID: {new_execution.execution_id} for message_id: {message_id}", level=logging.INFO)
        except Exception as db_error:
            db.rollback()
            logger.log_message(f"Error saving code execution: {str(db_error)}", level=logging.ERROR)
        finally:
            db.close()
        
        # Format plotly outputs for frontend
        plotly_outputs = [f"```plotly\n{json_output}\n```\n" for json_output in json_outputs]
        
        # Format matplotlib outputs for frontend
        matplotlib_chart_outputs = [f"```matplotlib\n{img_base64}\n```\n" for img_base64 in matplotlib_outputs]
        
        # Include execution status in the response
        return {
            "output": full_output,
            "plotly_outputs": plotly_outputs if json_outputs else None,
            "matplotlib_outputs": matplotlib_chart_outputs if matplotlib_outputs else None,
            "is_successful": is_successful,
            "failed_agents": failed_agents
        }
    except Exception as e:
        logger.log_message(f"Error executing code: {str(e)}", level=logging.ERROR)
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/edit")
async def edit_code(
    request_data: CodeEditRequest,
    request: Request,
    session_id: str = Depends(get_session_id_dependency)
):
    """
    Edit code provided in the request using AI
    
    Args:
        request_data: Body containing original code and user prompt
        request: FastAPI Request object
        session_id: Session identifier
        
    Returns:
        Dictionary containing the edited code
    """
    try:
        # Check if code and prompt are provided
        if not request_data.original_code or not request_data.user_prompt:
            raise HTTPException(status_code=400, detail="Both original code and editing instructions are required")
            
        # Access app state via request
        app_state = request.app.state
        session_state = app_state.get_session_state(session_id)
        
        # Get dataset context
        dataset_context = get_dataset_context(session_state["datasets"])
        try:
            # Use the configured language model with dataset context
            edited_code = edit_code_with_dspy(
                request_data.original_code, 
                request_data.user_prompt,
                dataset_context
            )
            edited_code = format_code_block(edited_code)
            return {
                "edited_code": edited_code,
            }
        except Exception as e:
            # Fallback if DSPy models are not initialized or there's an error
            logger.log_message(f"Error with DSPy models: {str(e)}", level=logging.ERROR)
            
            # Return a helpful error message that doesn't expose implementation details
            return {
                "edited_code": request_data.original_code,
                "error": "Could not process edit request. Please try again later."
            }
    except Exception as e:
        logger.log_message(f"Error editing code: {str(e)}", level=logging.ERROR)
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/fix")
async def fix_code(
    request_data: CodeFixRequest,
    request: Request,
    session_id: str = Depends(get_session_id_dependency)
):
    """
    Fix code with errors using block-by-block approach
    
    Args:
        request_data: Body containing code and error message
        request: FastAPI Request object
        session_id: Session identifier
        
    Returns:
        Dictionary containing the fixed code and information about fixed blocks
    """
    try:
        # Add debugging at the start
        logger.log_message(f"🔧 /fix endpoint called with session_id: {session_id}", level=logging.INFO)
        logger.log_message(f"🔧 Code length: {len(request_data.code) if request_data.code else 0}", level=logging.INFO)
        logger.log_message(f"🔧 Error length: {len(request_data.error) if request_data.error else 0}", level=logging.INFO)
        
        # Check if code and error are provided
        if not request_data.code or not request_data.error:
            logger.log_message(f"Error fixing code: Both code and error message are required {request_data.code} {request_data.error}", level=logging.ERROR)
            raise HTTPException(status_code=400, detail="Both code and error message are required")
            
        # Access app state via request
        app_state = request.app.state
        session_state = app_state.get_session_state(session_id)
        
        logger.log_message(f"🔧 Session state keys: {list(session_state.keys()) if session_state else 'None'}", level=logging.INFO)
        
        # Get the user_id from session state if available (for logging/tracking)
        user_id = session_state.get("user_id")
        logger.log_message(f"Code fix request from user_id: {user_id}, session_id: {session_id}", level=logging.INFO)
        
        # Get dataset context
        logger.log_message(f"🔧 Getting dataset context...", level=logging.INFO)
        dataset_context = get_dataset_context(session_state["datasets"])
        logger.log_message(f"🔧 Dataset context length: {len(dataset_context)}", level=logging.INFO)
        
        try:
            logger.log_message(f"🔧 Calling fix_code_with_dspy...", level=logging.INFO)
            # Use the code_fix agent to fix the code, with dataset context
            fixed_code = await fix_code_with_dspy(
                request_data.code, 
                request_data.error,
                dataset_context,
                session_state["datasets"]  # Pass the actual datasets
            )
            
            logger.log_message(f"🔧 fix_code_with_dspy returned, formatting...", level=logging.INFO)
            fixed_code = format_code_block(fixed_code)
            
            logger.log_message(f"Code fix completed successfully for user_id: {user_id}", level=logging.INFO)
            logger.log_message(f"🔧 Fixed code length: {len(fixed_code)}", level=logging.INFO)
                
            return {
                "fixed_code": fixed_code,
            }
        except Exception as e:
            logger.log_message(f"🔧 Error in fix_code_with_dspy: {str(e)}", level=logging.ERROR)
            # Fallback if DSPy models are not initialized or there's an error
            logger.log_message(f"Error with DSPy models for user_id {user_id}: {str(e)}", level=logging.ERROR)
            
            # Return the actual error details instead of generic message
            error_message = str(e)
            
            # Sanitize sensitive information but keep useful details
            if "API key" in error_message.lower():
                error_message = "API configuration error. Please contact support."
            elif "timeout" in error_message.lower():
                error_message = "Request timed out. Please try again."
            elif "rate limit" in error_message.lower():
                error_message = "Rate limit exceeded. Please wait a moment and try again."
            elif len(error_message) > 200:
                # Truncate very long error messages
                error_message = error_message[:200] + "..."
            
            return {
                "fixed_code": request_data.code,
                "error": error_message  # Return actual error instead of generic message
            }
    except Exception as e:
        logger.log_message(f"Error fixing code: {str(e)}", level=logging.ERROR)
        raise HTTPException(status_code=500, detail=str(e))
    
    
@router.post("/clean-code")
async def clean_code(
    request_data: CodeCleanRequest,
    request: Request,
    session_id: str = Depends(get_session_id_dependency)
):
    """
    Clean code provided in the request
    
    Args:
        request_data: Body containing code to clean
        request: FastAPI Request object
        session_id: Session identifier
        
    Returns:
        Dictionary containing the cleaned code
    """
    try:
        # Check if code is provided
        if not request_data.code:
            raise HTTPException(status_code=400, detail="Code is required")

        # Clean the code using the format_code function
        cleaned = format_code(request_data.code)
        
        return {
            "cleaned_code": cleaned,
        }
    except Exception as e:
        logger.log_message(f"Error cleaning code: {str(e)}", level=logging.ERROR)
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/get-latest-code")
async def get_latest_code(
    request_data: GetLatestCodeRequest,
    request: Request,
    session_id: str = Depends(get_session_id_dependency)
):
    """
    Retrieve the latest code for a specific message_id
    
    Args:
        request_data: Body containing message_id
        request: FastAPI Request object
        session_id: Session identifier
        
    Returns:
        Dictionary containing the latest code and execution status
    """
    try:
        message_id = request_data.message_id
        
        if not message_id:
            raise HTTPException(status_code=400, detail="Message ID is required")
            
        # Get database session
        db = get_session()
        
        try:
            # Query the database for the latest code execution record (ordered by created_at desc)
            logger.log_message(f"Searching for execution records with message_id: {message_id}", level=logging.INFO)
            
            execution_record = db.query(CodeExecution).filter(
                CodeExecution.message_id == message_id
            ).order_by(CodeExecution.created_at.desc()).first()
            
            # Also log total count of records for this message_id
            total_records = db.query(CodeExecution).filter(
                CodeExecution.message_id == message_id
            ).count()
            logger.log_message(f"Found {total_records} execution records for message_id: {message_id}", level=logging.INFO)
            
            if execution_record:
                logger.log_message(f"Latest execution record found - success: {execution_record.is_successful}, latest_code length: {len(execution_record.latest_code or '') if execution_record.latest_code else 0}", level=logging.INFO)

                # Return the latest code and execution status
                return {
                    "found": True,
                    "message_id": message_id,
                    "latest_code": execution_record.latest_code,
                    "initial_code": execution_record.initial_code,
                    "is_successful": execution_record.is_successful,
                    "failed_agents": execution_record.failed_agents
                }
            else:
                logger.log_message(f"No execution record found for message_id: {message_id}", level=logging.INFO)
                return {
                    "found": False,
                    "message_id": message_id
                }
                
        except Exception as db_error:
            logger.log_message(f"Database error retrieving latest code: {str(db_error)}", level=logging.ERROR)
            raise HTTPException(status_code=500, detail=f"Database error: {str(db_error)}")
        finally:
            db.close()
            
    except Exception as e:
        logger.log_message(f"Error retrieving latest code: {str(e)}", level=logging.ERROR)
        raise HTTPException(status_code=500, detail=str(e))