import base64
import os
import glob
from openai import OpenAI

from gpt_agent import kimi_api, openai_api, deep_seek_api, generate


def observe_task_execution(episode_id, task_name, task_info, problematic_code=None, save_dir="./camera_images", camera_name=None, generate_dir_name=None):
    """
    Observe task execution by analyzing step-by-step images using an image understanding API.

    Args:
    episode_id (int): ID of the episode to analyze.
    task_name (str): Name of the task.
    task_info (dict): Basic information about the task.
    problematic_code (str, optional): Potentially faulty code generated in a previous step.
    save_dir (str): Base directory where images are saved.
    camera_name (str): Name of the camera used to capture the images.
    generate_dir_name (str, optional): Name of the subdirectory with generated images.

    Returns:
    str: Textual description of the observation result.
    """
    client = OpenAI(
        api_key=kimi_api,
        base_url="https://api.moonshot.cn/v1",
    )
    
    # Check if the save_dir already contains the task name
    base_task_name = task_name.lower() if task_name else ""
    if base_task_name and os.path.basename(save_dir) == base_task_name:
        # If task name is already included in save_dir, use it directly
        task_dir = save_dir
    else:
        # Otherwise, append task name to the path
        task_dir = os.path.join(save_dir, base_task_name) if base_task_name else save_dir
    
    # If a generated subdirectory name is specified, add it to the path
    if generate_dir_name:
        task_dir = os.path.join(task_dir, generate_dir_name)
    
    print(f"Looking for task images in: {os.path.abspath(task_dir)}")
    
    # Check if task directory exists
    if not os.path.exists(task_dir):
        return f"Error: Image directory not found at {task_dir}"
    
    # Get images for the specific episode
    image_files = sorted(glob.glob(os.path.join(task_dir, f"episode{episode_id}_*.png")))
    
    if not image_files:
        return f"Error: No images found for episode {episode_id} in directory {task_dir}"
    
    # Extract step names from image filenames
    step_names = []
    for f in image_files:
        filename = os.path.basename(f)
        first_underscore_pos = filename.find('_')
        if first_underscore_pos != -1:
            step_name = filename[first_underscore_pos+1:].rsplit('.', 1)[0]
            step_names.append(step_name)
        else:
            step_names.append(filename.rsplit('.', 1)[0])

    # Logging for debugging purposes (from observation_agent.py)
    print(f"Image search pattern: episode{episode_id}_*.png, number of files found: {len(image_files)}")
    # for f in image_files[:5]:  # Uncomment to print first 5 filenames
    #     print(f"  - {os.path.basename(f)}")
    
    # Construct the prompt
    prompt = f"""Analyze the execution of the following robot task:
Task name: {task_name}
Task description: {task_info.get('description', 'No description provided')}
Task goal: {task_info.get('goal', 'No goal provided')}

You will be shown images from each step of the task execution. Please analyze:
1. Whether each step was executed successfully.
2. If any step failed, identify which one and explain why.
3. Whether the overall task was successfully completed.
4. If the task failed, provide detailed reasoning.

You will see execution images for the following steps: {', '.join(step_names)}
"""

    if problematic_code:
        prompt += f"\nHere is a piece of potentially problematic code:\n```python\n{problematic_code}\n```\nPlease analyze if the code is related to the observed issue."
    
    # Prepare message content for API call
    user_content = []
    
    # Add textual prompt
    user_content.append({
        "type": "text",
        "text": prompt
    })
    
    # Add images and step names
    for img_path in image_files:
        filename = os.path.basename(img_path)
        first_underscore_pos = filename.find('_')
        if first_underscore_pos != -1:
            step_name = filename[first_underscore_pos+1:].rsplit('.', 1)[0]
        else:
            step_name = filename.rsplit('.', 1)[0]
        
        # Add step name
        user_content.append({
            "type": "text",
            "text": f"Step: {step_name}"
        })
        
        # Add image as base64
        try:
            base64_image = encode_image(img_path)
            user_content.append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/png;base64,{base64_image}"
                }
            })
        except Exception as e:
            print(f"Warning: Failed to encode image {img_path}: {str(e)}")
    
    # Call the image analysis API
    try:
        response = client.chat.completions.create(
            model="moonshot-v1-32k-vision-preview",
            messages=[
                {"role": "system", "content": "You are a robot task execution analysis expert. Please analyze the provided image sequence."},
                {"role": "user", "content": user_content}
            ]
        )
        return response.choices[0].message.content
    except Exception as e:
        error_msg = f"Error occurred while calling the image understanding API: {str(e)}"
        print(error_msg)
        return error_msg


def encode_image(image_path):
    """Encode an image file to a base64 string."""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


def insert_observation_points(task_info, task_code, generate_num_id=0):
    """
    Insert observation function calls at key points in robot task code.
    
    Args:
        task_info (dict): Information about the task
        task_code (str): Original code for the task
        
    Returns:
        str: Code with inserted observation points and steps summary
    """

    # Extract task name
    if isinstance(task_info, dict) and 'task_name' in task_info:
        task_name = task_info.get('task_name')
    else:
        # Try to extract from code
        import re
        task_name_match = re.search(r'class\s+gpt_(\w+)', task_code)
        task_name = task_name_match.group(1) if task_name_match else "unknown_task"

    # Prepare the prompt for the LLM
    prompt = f"""You are an expert in robot programming. I have a robot task code that needs observation functions added for monitoring.

Task information:
{task_info}

I need you to:
1. Identify ONLY the main logical steps in this task implementation that cause SIGNIFICANT SCENE CHANGES
2. After each such logical step in the code, insert a camera observation function with this format:
   `self.save_camera_images(task_name="{task_name}", step_name="stepX_descriptive_name", generate_num_id="generate_num_{generate_num_id}")`
   where X is the sequential step number and descriptive_name is a brief description of what just happened
3. Provide a numbered list of all the steps you've identified in the task
4. ADD AN OBSERVATION AT THE BEGINNING OF THE TASK to capture the initial scene state
5. ADD AN OBSERVATION AT THE END OF THE TASK to capture the final scene state

Here's the current code:
```python
{task_code}
```

IMPORTANT CONSTRAINTS:
- ADD FEWER THAN 10 OBSERVATION POINTS in total
- ONLY add observations after operations that cause VISIBLE SCENE CHANGES
- Do NOT add observations for planning, calculations, or any operations that don't visibly change the scene
- Focus on key state changes like: robot arm movements, gripper operations, object manipulations
- Skip observations for intermediate movements, planning steps, or calculations
- The observation function is already defined in the code
- Give each step a descriptive name like "gripper_closed", "move_to_target", etc.
- The step number (X in stepX) should increase sequentially
- DO NOT MODIFY ANY EXISTING ROBOT OPERATION CODE - only insert observation function calls after existing code without changing the original functionality

Format your response as follows:

STEP_LIST:
1. First step description
2. Second step description
...

MODIFIED_CODE:
```python
<the entire modified code with observation functions inserted>
```
"""
    
    # Get the modified code from LLM in one call
    response = generate(message=[{
        "role": "system",
        "content": "You are an AI assistant that helps with programming robot tasks."
    }, {
        "role": "user",
        "content": prompt
    }])

    # Extract the step list and modified code
    try:
        steps_part, code_part = response.split("MODIFIED_CODE:", 1)
        steps = steps_part.replace("STEP_LIST:", "").strip()
        modified_code = code_part.strip()

        # Clean up any potential markdown code block formatting
        if modified_code.startswith("```python"):
            modified_code = modified_code[len("```python"):].strip()
        if modified_code.endswith("```"):
            modified_code = modified_code[:-3].strip()
    except ValueError:
        # Fallback in case the format isn't as expected
        steps = "Failed to extract step list"
        modified_code = response

    # Format the output
    output = f"# task_name: {task_name}\n# task_step:\n{steps}\n\n# task_code:\n```python\n{modified_code}\n```"

    return output
