"""
Trajectory saving callback handler for ComputerAgent.
"""

import base64
import io
import json
import os
import uuid
from copy import deepcopy
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

try:
    from typing import override
except ImportError:
    from typing_extensions import override

from PIL import Image, ImageDraw

from .base import AsyncCallbackHandler


def sanitize_image_urls(data: Any) -> Any:
    """
    Recursively search for 'image_url' keys and set their values to '[omitted]'.

    Args:
        data: Any data structure (dict, list, or primitive type)

    Returns:
        A deep copy of the data with all 'image_url' values replaced with '[omitted]'
    """
    if isinstance(data, dict):
        # Create a copy of the dictionary
        sanitized = {}
        for key, value in data.items():
            if key == "image_url":
                sanitized[key] = "[omitted]"
            else:
                # Recursively sanitize the value
                sanitized[key] = sanitize_image_urls(value)
        return sanitized

    elif isinstance(data, list):
        # Recursively sanitize each item in the list
        return [sanitize_image_urls(item) for item in data]

    else:
        # For primitive types (str, int, bool, None, etc.), return as-is
        return data


def extract_computer_call_outputs(
    items: List[Dict[str, Any]], screenshot_dir: Optional[Path]
) -> List[Dict[str, Any]]:
    """
    Save any base64-encoded screenshots from computer_call_output or function_call_output
    entries to files and replace their image_url with the saved file path when a call_id is present.

    Only operates if screenshot_dir is provided and exists; otherwise returns items unchanged.

    Args:
        items: List of message/result dicts potentially containing computer_call_output
               or function_call_output entries
        screenshot_dir: Directory to write screenshots into

    Returns:
        A new list with updated image_url fields when applicable.
    """
    if not items:
        return items
    if not screenshot_dir or not screenshot_dir.exists():
        return items

    updated: List[Dict[str, Any]] = []
    for item in items:
        # work on a shallow copy; deep copy nested 'output' if we modify it
        msg = dict(item)
        try:
            if msg.get("type") == "computer_call_output":
                call_id = msg.get("call_id")
                output = msg.get("output", {})
                image_url = output.get("image_url")
                if call_id and isinstance(image_url, str) and image_url.startswith("data:"):
                    # derive extension from MIME type e.g. data:image/png;base64,
                    try:
                        ext = image_url.split(";", 1)[0].split("/")[-1]
                        if not ext:
                            ext = "png"
                    except Exception:
                        ext = "png"
                    out_path = screenshot_dir / f"{call_id}.{ext}"
                    # write file if it doesn't exist
                    if not out_path.exists():
                        try:
                            b64_payload = image_url.split(",", 1)[1]
                            img_bytes = base64.b64decode(b64_payload)
                            out_path.parent.mkdir(parents=True, exist_ok=True)
                            with open(out_path, "wb") as f:
                                f.write(img_bytes)
                        except Exception:
                            # if anything fails, skip modifying this message
                            pass
                    # update image_url to file path
                    new_output = dict(output)
                    new_output["image_url"] = str(out_path)
                    msg["output"] = new_output

            elif msg.get("type") == "function_call_output":
                # Handle function_call_output from GPT 5.4 / BrowserTool
                call_id = msg.get("call_id")
                output = msg.get("output", "")

                # Parse output if it's a string
                if isinstance(output, str):
                    try:
                        output_dict = json.loads(output)
                    except (json.JSONDecodeError, TypeError):
                        output_dict = None
                else:
                    output_dict = output

                if isinstance(output_dict, dict) and call_id:
                    image_data = None
                    image_key = None

                    # Format 1: {"type": "input_image", "image_url": "data:image/png;base64,..."}
                    if output_dict.get("type") == "input_image":
                        image_url = output_dict.get("image_url", "")
                        if isinstance(image_url, str) and image_url.startswith("data:"):
                            image_data = image_url.split(",", 1)[1] if "," in image_url else None
                            image_key = "image_url"

                    # Format 2: {"success": True, "screenshot": "base64data"}
                    elif output_dict.get("screenshot"):
                        image_data = output_dict.get("screenshot")
                        image_key = "screenshot"

                    if image_data and image_key:
                        out_path = screenshot_dir / f"{call_id}.png"
                        if not out_path.exists():
                            try:
                                img_bytes = base64.b64decode(image_data)
                                out_path.parent.mkdir(parents=True, exist_ok=True)
                                with open(out_path, "wb") as f:
                                    f.write(img_bytes)
                            except Exception:
                                pass
                        # Update output to reference file path
                        new_output_dict = dict(output_dict)
                        new_output_dict[image_key] = str(out_path)
                        msg["output"] = json.dumps(new_output_dict)

            elif msg.get("role") == "user":
                # Handle user messages with input_image content (GPT-5.4 sibling screenshot messages)
                # These accompany function_call_output for computer calls
                content = msg.get("content", [])
                if isinstance(content, list):
                    new_content = []
                    content_modified = False
                    for content_item in content:
                        if (
                            isinstance(content_item, dict)
                            and content_item.get("type") == "input_image"
                        ):
                            image_url = content_item.get("image_url", "")
                            if isinstance(image_url, str) and image_url.startswith("data:"):
                                # Generate a unique ID for this screenshot
                                screenshot_id = str(uuid.uuid4())[:8]
                                try:
                                    ext = image_url.split(";", 1)[0].split("/")[-1]
                                    if not ext:
                                        ext = "png"
                                except Exception:
                                    ext = "png"
                                out_path = screenshot_dir / f"user_screenshot_{screenshot_id}.{ext}"
                                if not out_path.exists():
                                    try:
                                        b64_payload = image_url.split(",", 1)[1]
                                        img_bytes = base64.b64decode(b64_payload)
                                        out_path.parent.mkdir(parents=True, exist_ok=True)
                                        with open(out_path, "wb") as f:
                                            f.write(img_bytes)
                                    except Exception:
                                        new_content.append(content_item)
                                        continue
                                # Update image_url to file path
                                new_item = dict(content_item)
                                new_item["image_url"] = str(out_path)
                                new_content.append(new_item)
                                content_modified = True
                            else:
                                new_content.append(content_item)
                        else:
                            new_content.append(content_item)
                    if content_modified:
                        msg["content"] = new_content

        except Exception:
            # do not block on malformed entries; keep original
            pass
        updated.append(msg)
    return updated


class TrajectorySaverCallback(AsyncCallbackHandler):
    """
    Callback handler that saves agent trajectories to disk.

    Saves each run as a separate trajectory with unique ID, and each turn
    within the trajectory gets its own folder with screenshots and responses.
    """

    def __init__(
        self, trajectory_dir: str, reset_on_run: bool = True, screenshot_dir: Optional[str] = None
    ):
        """
        Initialize trajectory saver.

        Args:
            trajectory_dir: Base directory to save trajectories
            reset_on_run: If True, reset trajectory_id/turn/artifact on each run.
                         If False, continue using existing trajectory_id if set.
        """
        self.trajectory_dir = Path(trajectory_dir)
        self.trajectory_id: Optional[str] = None
        self.current_turn: int = 0
        self.current_artifact: int = 0
        self.model: Optional[str] = None
        self.total_usage: Dict[str, Any] = {}
        self.reset_on_run = reset_on_run
        # Optional directory to store extracted screenshots from metadata/new_items
        self.screenshot_dir: Optional[Path] = Path(screenshot_dir) if screenshot_dir else None

        # Ensure trajectory directory exists
        self.trajectory_dir.mkdir(parents=True, exist_ok=True)

        # Ensure screenshot directory exists if specified
        if self.screenshot_dir:
            self.screenshot_dir.mkdir(parents=True, exist_ok=True)

    def _get_turn_dir(self) -> Path:
        """Get the directory for the current turn."""
        if not self.trajectory_id:
            raise ValueError("Trajectory not initialized - call _on_run_start first")

        # format: trajectory_id/turn_000
        turn_dir = self.trajectory_dir / self.trajectory_id / f"turn_{self.current_turn:03d}"
        turn_dir.mkdir(parents=True, exist_ok=True)
        return turn_dir

    def _save_artifact(self, name: str, artifact: Union[str, bytes, Dict[str, Any]]) -> None:
        """Save an artifact to the current turn directory."""
        turn_dir = self._get_turn_dir()
        if isinstance(artifact, bytes):
            # format: turn_000/0000_name.png
            artifact_filename = f"{self.current_artifact:04d}_{name}"
            artifact_path = turn_dir / f"{artifact_filename}.png"
            with open(artifact_path, "wb") as f:
                f.write(artifact)
        else:
            # format: turn_000/0000_name.json
            artifact_filename = f"{self.current_artifact:04d}_{name}"
            artifact_path = turn_dir / f"{artifact_filename}.json"
            # add created_at
            if isinstance(artifact, dict):
                artifact = artifact.copy()
                artifact["created_at"] = str(uuid.uuid1().time)
            with open(artifact_path, "w") as f:
                json.dump(sanitize_image_urls(artifact), f, indent=2)
        self.current_artifact += 1

    def _update_usage(self, usage: Dict[str, Any]) -> None:
        """Update total usage statistics."""

        def add_dicts(target: Dict[str, Any], source: Dict[str, Any]) -> None:
            for key, value in source.items():
                if isinstance(value, dict):
                    if key not in target:
                        target[key] = {}
                    add_dicts(target[key], value)
                else:
                    if key not in target:
                        target[key] = 0
                    target[key] += value

        add_dicts(self.total_usage, usage)

    @override
    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
        """Initialize trajectory tracking for a new run."""
        model = kwargs.get("model", "unknown")

        # Only reset trajectory state if reset_on_run is True or no trajectory exists
        if self.reset_on_run or not self.trajectory_id:
            model_name_short = model.split("+")[-1].split("/")[-1].lower()[:16]
            if "+" in model:
                model_name_short = model.split("+")[0].lower()[:4] + "_" + model_name_short
            # strip non-alphanumeric characters from model_name_short
            model_name_short = "".join(c for c in model_name_short if c.isalnum() or c == "_")

            # id format: yyyy-mm-dd_model_hhmmss_uuid[:4]
            now = datetime.now()
            self.trajectory_id = f"{now.strftime('%Y-%m-%d')}_{model_name_short}_{now.strftime('%H%M%S')}_{str(uuid.uuid4())[:4]}"
            self.current_turn = 0
            self.current_artifact = 0
            self.model = model
            self.total_usage = {}

            # Create trajectory directory
            trajectory_path = self.trajectory_dir / self.trajectory_id
            trajectory_path.mkdir(parents=True, exist_ok=True)

            # Save trajectory metadata (optionally extract screenshots to screenshot_dir)
            kwargs_to_save = kwargs.copy()
            try:
                if "messages" in kwargs_to_save:
                    kwargs_to_save["messages"] = extract_computer_call_outputs(
                        kwargs_to_save["messages"], self.screenshot_dir
                    )
            except Exception:
                # If extraction fails, fall back to original messages
                pass
            metadata = {
                "trajectory_id": self.trajectory_id,
                "created_at": str(uuid.uuid1().time),
                "status": "running",
                "kwargs": kwargs_to_save,
            }

            with open(trajectory_path / "metadata.json", "w") as f:
                json.dump(metadata, f, indent=2)
        else:
            # Continue with existing trajectory - just update model if needed
            self.model = model

    @override
    async def on_run_end(
        self,
        kwargs: Dict[str, Any],
        old_items: List[Dict[str, Any]],
        new_items: List[Dict[str, Any]],
    ) -> None:
        """Finalize run tracking by updating metadata with completion status, usage, and new items."""
        if not self.trajectory_id:
            return

        # Update metadata with completion status, total usage, and new items
        trajectory_path = self.trajectory_dir / self.trajectory_id
        metadata_path = trajectory_path / "metadata.json"

        # Read existing metadata
        if metadata_path.exists():
            with open(metadata_path, "r") as f:
                metadata = json.load(f)
        else:
            metadata = {}

        # Update metadata with completion info
        # Optionally extract screenshots from new_items before persisting
        new_items_to_save = new_items
        try:
            new_items_to_save = extract_computer_call_outputs(new_items, self.screenshot_dir)
        except Exception:
            pass

        metadata.update(
            {
                "status": "completed",
                "completed_at": str(uuid.uuid1().time),
                "total_usage": self.total_usage,
                "new_items": new_items_to_save,
                "total_turns": self.current_turn,
            }
        )

        # Save updated metadata
        with open(metadata_path, "w") as f:
            json.dump(metadata, f, indent=2)

    @override
    async def on_api_start(self, kwargs: Dict[str, Any]) -> None:
        if not self.trajectory_id:
            return

        self._save_artifact("api_start", {"kwargs": kwargs})

    @override
    async def on_api_end(self, kwargs: Dict[str, Any], result: Any) -> None:
        """Save API call result."""
        if not self.trajectory_id:
            return

        self._save_artifact("api_result", {"kwargs": kwargs, "result": result})

    @override
    async def on_screenshot(self, screenshot: Union[str, bytes], name: str = "screenshot") -> None:
        """Save a screenshot."""
        if isinstance(screenshot, str):
            screenshot = base64.b64decode(screenshot)
        self._save_artifact(name, screenshot)

    @override
    async def on_usage(self, usage: Dict[str, Any]) -> None:
        """Called when usage information is received."""
        self._update_usage(usage)

    @override
    async def on_responses(self, kwargs: Dict[str, Any], responses: Dict[str, Any]) -> None:
        """Save responses to the current turn directory and update usage statistics."""
        if not self.trajectory_id:
            return

        # Save responses
        turn_dir = self._get_turn_dir()
        response_data = {
            "timestamp": str(uuid.uuid1().time),
            "model": self.model,
            "kwargs": kwargs,
            "response": responses,
        }

        self._save_artifact("agent_response", response_data)

        # Increment turn counter
        self.current_turn += 1

    def _draw_crosshair_on_image(self, image_bytes: bytes, x: int, y: int) -> bytes:
        """
        Draw a red dot and crosshair at the specified coordinates on the image.

        Args:
            image_bytes: The original image as bytes
            x: X coordinate for the crosshair
            y: Y coordinate for the crosshair

        Returns:
            Modified image as bytes with red dot and crosshair
        """
        # Open the image
        image = Image.open(io.BytesIO(image_bytes))
        draw = ImageDraw.Draw(image)

        # Draw crosshair lines (red, 2px thick)
        crosshair_size = 20
        line_width = 2
        color = "red"

        # Horizontal line
        draw.line([(x - crosshair_size, y), (x + crosshair_size, y)], fill=color, width=line_width)
        # Vertical line
        draw.line([(x, y - crosshair_size), (x, y + crosshair_size)], fill=color, width=line_width)

        # Draw center dot (filled circle)
        dot_radius = 3
        draw.ellipse(
            [(x - dot_radius, y - dot_radius), (x + dot_radius, y + dot_radius)], fill=color
        )

        # Convert back to bytes
        output = io.BytesIO()
        image.save(output, format="PNG")
        return output.getvalue()

    @override
    async def on_computer_call_end(
        self, item: Dict[str, Any], result: List[Dict[str, Any]]
    ) -> None:
        """
        Called when a computer call has completed.
        Saves screenshots and computer call output.
        """
        if not self.trajectory_id:
            return

        self._save_artifact("computer_call_result", {"item": item, "result": result})

        # Check if action has x/y coordinates and there's a screenshot in the result
        action = item.get("action", {})
        if "x" in action and "y" in action:
            # Look for screenshot in the result
            for result_item in result:
                if (
                    result_item.get("type") == "computer_call_output"
                    and result_item.get("output", {}).get("type") == "input_image"
                ):

                    image_url = result_item["output"]["image_url"]

                    # Extract base64 image data
                    if image_url.startswith("data:image/"):
                        # Format: data:image/png;base64,<base64_data>
                        base64_data = image_url.split(",", 1)[1]
                    else:
                        # Assume it's just base64 data
                        base64_data = image_url

                    try:
                        # Decode the image
                        image_bytes = base64.b64decode(base64_data)

                        # Draw crosshair at the action coordinates
                        annotated_image = self._draw_crosshair_on_image(
                            image_bytes, int(action["x"]), int(action["y"])
                        )

                        # Save as screenshot_action
                        self._save_artifact("screenshot_action", annotated_image)

                    except Exception as e:
                        # If annotation fails, just log and continue
                        print(f"Failed to annotate screenshot: {e}")

                    break  # Only process the first screenshot found

        # Increment turn counter
        self.current_turn += 1

    @override
    async def on_function_call_end(
        self, item: Dict[str, Any], result: List[Dict[str, Any]]
    ) -> None:
        """
        Called when a function call has completed.
        Saves screenshots and function call output for GPT 5.4 / BrowserTool.
        """
        if not self.trajectory_id:
            return

        self._save_artifact("function_call_result", {"item": item, "result": result})

        # Extract coordinates from function call arguments if present
        x_coord, y_coord = None, None
        try:
            arguments = item.get("arguments", "{}")
            if isinstance(arguments, str):
                args_dict = json.loads(arguments)
            else:
                args_dict = arguments

            # Check for coordinate array format (BrowserTool style)
            coord = args_dict.get("coordinate")
            if coord and isinstance(coord, list) and len(coord) >= 2:
                x_coord, y_coord = coord[0], coord[1]
            # Check for x/y format (computer_use style)
            elif "x" in args_dict and "y" in args_dict:
                x_coord, y_coord = args_dict.get("x"), args_dict.get("y")
        except (json.JSONDecodeError, TypeError):
            pass

        # Look for screenshot in the result
        screenshot_found = False
        for result_item in result:
            if screenshot_found:
                break

            if result_item.get("type") == "function_call_output":
                output = result_item.get("output", "")

                # Parse output if it's a string
                if isinstance(output, str):
                    try:
                        output_dict = json.loads(output)
                    except (json.JSONDecodeError, TypeError):
                        # Try to evaluate as Python literal (for stringified dicts)
                        try:
                            import ast

                            output_dict = ast.literal_eval(output)
                        except (ValueError, SyntaxError):
                            continue
                else:
                    output_dict = output

                if not isinstance(output_dict, dict):
                    continue

                # Extract screenshot from various formats
                image_data = None

                # Format 1: {"type": "input_image", "image_url": "data:image/png;base64,..."}
                if output_dict.get("type") == "input_image":
                    image_url = output_dict.get("image_url", "")
                    if image_url.startswith("data:image/"):
                        image_data = image_url.split(",", 1)[1]
                    elif image_url:
                        image_data = image_url

                # Format 2: {"success": True, "screenshot": "base64data"}
                elif output_dict.get("screenshot"):
                    image_data = output_dict.get("screenshot")

                if image_data:
                    try:
                        # Decode the image
                        image_bytes = base64.b64decode(image_data)

                        # If we have coordinates, draw crosshair annotation
                        if (
                            x_coord is not None
                            and y_coord is not None
                            and x_coord != 0
                            and y_coord != 0
                        ):
                            annotated_image = self._draw_crosshair_on_image(
                                image_bytes, int(x_coord), int(y_coord)
                            )
                            self._save_artifact("screenshot_action", annotated_image)
                        else:
                            # Save plain screenshot without crosshair
                            self._save_artifact("screenshot", image_bytes)

                        screenshot_found = True

                    except Exception as e:
                        # If processing fails, just log and continue
                        print(f"Failed to process screenshot from function call: {e}")

            # Handle sibling user messages with input_image content (GPT-5.4 computer calls)
            # These accompany function_call_output and contain the actual screenshot
            elif result_item.get("role") == "user":
                content = result_item.get("content", [])
                if isinstance(content, list):
                    for content_item in content:
                        if (
                            isinstance(content_item, dict)
                            and content_item.get("type") == "input_image"
                        ):
                            image_url = content_item.get("image_url", "")
                            if isinstance(image_url, str) and image_url.startswith("data:"):
                                try:
                                    b64_payload = image_url.split(",", 1)[1]
                                    image_bytes = base64.b64decode(b64_payload)

                                    # If we have coordinates, draw crosshair annotation
                                    if (
                                        x_coord is not None
                                        and y_coord is not None
                                        and x_coord != 0
                                        and y_coord != 0
                                    ):
                                        annotated_image = self._draw_crosshair_on_image(
                                            image_bytes, int(x_coord), int(y_coord)
                                        )
                                        self._save_artifact("screenshot_action", annotated_image)
                                    else:
                                        # Save plain screenshot without crosshair
                                        self._save_artifact("screenshot", image_bytes)

                                    screenshot_found = True
                                    break

                                except Exception as e:
                                    # If processing fails, just log and continue
                                    print(f"Failed to process screenshot from user message: {e}")

        # Increment turn counter
        self.current_turn += 1