import base64
import json
import os
import shlex
from pathlib import Path
from typing import Any, Literal

from harbor.agents.installed.base import BaseInstalledAgent, CliFlag, ExecInput
from harbor.models.agent.context import AgentContext
from harbor.models.agent.name import AgentName
from harbor.models.trajectories import (
    Agent,
    ContentPart,
    FinalMetrics,
    ImageSource,
    Metrics,
    Observation,
    ObservationResult,
    Step,
    ToolCall,
    Trajectory,
)

_ImageMediaType = Literal["image/jpeg", "image/png", "image/gif", "image/webp"]


class GeminiCli(BaseInstalledAgent):
    """
    The Gemini CLI agent uses Google's Gemini CLI tool to solve tasks.
    """

    def get_version_command(self) -> str | None:
        return ". ~/.nvm/nvm.sh; gemini --version"

    SUPPORTS_ATIF: bool = True

    CLI_FLAGS = [
        CliFlag(
            "sandbox",
            cli="--sandbox",
            type="bool",
        ),
    ]

    # Counter for generating unique image filenames within a session
    _image_counter: int = 0

    @staticmethod
    def name() -> str:
        return AgentName.GEMINI_CLI.value

    @property
    def _install_agent_template_path(self) -> Path:
        return Path(__file__).parent / "install-gemini-cli.sh.j2"

    def _save_image(
        self,
        image_data: str,
        mime_type: str,
        step_id: int,
        obs_index: int,
        image_index: int = 0,
    ) -> tuple[str, _ImageMediaType] | tuple[None, None]:
        """Save a base64 image to the images directory.

        Args:
            image_data: Base64-encoded image data
            mime_type: MIME type of the image (e.g., 'image/png')
            step_id: The step ID this image belongs to
            obs_index: Index of the observation result within the step
            image_index: Index of the image within the observation (for multiple images)

        Returns:
            Tuple of (relative_path, media_type) for the saved image, or (None, None) on failure
        """
        # Create images directory if it doesn't exist
        images_dir = self.logs_dir / "images"
        images_dir.mkdir(exist_ok=True)

        # Determine file extension from mime type
        # Only accept MIME types that ImageSource validates
        extension_map: dict[_ImageMediaType, str] = {
            "image/png": "png",
            "image/jpeg": "jpg",
            "image/gif": "gif",
            "image/webp": "webp",
        }
        for valid_type, extension in extension_map.items():
            if mime_type == valid_type:
                break
        else:
            # Unsupported MIME type - return None to avoid Pydantic validation error
            print(f"Warning: Unsupported image MIME type: {mime_type}")
            return None, None

        # Generate unique filename
        filename = f"step_{step_id}_obs_{obs_index}_img_{image_index}.{extension}"
        image_path = images_dir / filename

        # Decode and save the image
        try:
            image_bytes = base64.b64decode(image_data)
            image_path.write_bytes(image_bytes)
        except Exception as e:
            print(f"Warning: Failed to save image: {e}")
            return None, None

        # Return relative path from trajectory.json location
        return f"images/{filename}", valid_type

    def _convert_gemini_to_atif(
        self, gemini_trajectory: dict[str, Any]
    ) -> Trajectory | None:
        """Convert Gemini CLI trajectory format to ATIF format."""
        session_id = gemini_trajectory.get("sessionId", "unknown")
        messages = gemini_trajectory.get("messages", [])

        if not messages:
            return None

        def _extract_text(content: Any) -> str:
            """Extract text from Gemini content field (list of dicts or string)."""
            if isinstance(content, str):
                return content
            if isinstance(content, list):
                return "\n".join(
                    part.get("text", "") if isinstance(part, dict) else str(part)
                    for part in content
                )
            return str(content) if content else ""

        steps: list[Step] = []
        step_id = 1

        # Track metrics for final_metrics calculation
        total_input_tokens = 0
        total_output_tokens = 0
        total_cached_tokens = 0

        for message in messages:
            msg_type = message.get("type")
            timestamp = message.get("timestamp")

            # User message
            if msg_type == "user":
                content = _extract_text(message.get("content", ""))
                steps.append(
                    Step(
                        step_id=step_id,
                        timestamp=timestamp,
                        source="user",
                        message=content,
                    )
                )
                step_id += 1

            # Gemini (agent) message
            elif msg_type == "gemini":
                content = _extract_text(message.get("content", ""))
                thoughts = message.get("thoughts", [])
                tool_calls_data = message.get("toolCalls", [])
                tokens = message.get("tokens", {})
                model_name = message.get("model")

                # Build reasoning content from thoughts
                reasoning_content: str | None = None
                if thoughts:
                    reasoning_parts = []
                    for thought in thoughts:
                        subject = thought.get("subject", "")
                        description = thought.get("description", "")
                        if subject and description:
                            reasoning_parts.append(f"{subject}: {description}")
                        elif description:
                            reasoning_parts.append(description)
                    if reasoning_parts:
                        reasoning_content = "\n".join(reasoning_parts)

                # Handle tool calls
                tool_calls: list[ToolCall] | None = None
                observation: Observation | None = None

                if tool_calls_data:
                    tool_calls = []
                    observation_results: list[ObservationResult] = []

                    for tc in tool_calls_data:
                        tool_call_id = tc.get("id", "")
                        tool_name = tc.get("name", "")
                        args = tc.get("args", {})
                        result = tc.get("result", [])

                        tool_calls.append(
                            ToolCall(
                                tool_call_id=tool_call_id,
                                function_name=tool_name,
                                arguments=args,
                            )
                        )

                        # Extract observation content from result
                        # This may include text output and/or image data
                        obs_content: str | list[ContentPart] | None = None
                        obs_index = len(observation_results)

                        if result:
                            text_output: str | None = None
                            image_parts: list[ContentPart] = []

                            for res_item in result:
                                if isinstance(res_item, dict):
                                    func_resp = res_item.get("functionResponse", {})
                                    response = func_resp.get("response", {})
                                    output = response.get("output")
                                    if output:
                                        text_output = output

                                    # Check for image data in parts
                                    parts = func_resp.get("parts", [])
                                    image_index = 0
                                    for part in parts:
                                        if isinstance(part, dict):
                                            inline_data = part.get("inlineData", {})
                                            if inline_data:
                                                mime_type = inline_data.get(
                                                    "mimeType", "image/png"
                                                )
                                                data = inline_data.get("data", "")
                                                if data:
                                                    # Save the image and get the path
                                                    image_path, media_type = (
                                                        self._save_image(
                                                            data,
                                                            mime_type,
                                                            step_id,
                                                            obs_index,
                                                            image_index,
                                                        )
                                                    )
                                                    if image_path and media_type:
                                                        image_parts.append(
                                                            ContentPart(
                                                                type="image",
                                                                source=ImageSource(
                                                                    media_type=media_type,
                                                                    path=image_path,
                                                                ),
                                                            )
                                                        )
                                                    image_index += 1

                            # Build observation content
                            if image_parts:
                                # Multimodal content - combine text and images
                                content_parts: list[ContentPart] = []
                                if text_output:
                                    content_parts.append(
                                        ContentPart(type="text", text=text_output)
                                    )
                                content_parts.extend(image_parts)
                                obs_content = content_parts
                            else:
                                # Text-only content
                                obs_content = text_output

                        observation_results.append(
                            ObservationResult(
                                source_call_id=tool_call_id or None,
                                content=obs_content,
                            )
                        )

                    if observation_results:
                        observation = Observation(results=observation_results)

                # Build metrics
                metrics: Metrics | None = None
                if tokens:
                    input_tokens = tokens.get("input", 0)
                    output_tokens = tokens.get("output", 0)
                    cached_tokens = tokens.get("cached", 0)
                    thoughts_tokens = tokens.get("thoughts", 0)
                    tool_tokens = tokens.get("tool", 0)

                    # Calculate completion tokens (output + thoughts + tool)
                    completion_tokens = output_tokens + thoughts_tokens + tool_tokens

                    # Update totals
                    total_input_tokens += input_tokens
                    total_output_tokens += completion_tokens
                    total_cached_tokens += cached_tokens

                    metrics = Metrics(
                        prompt_tokens=input_tokens,
                        completion_tokens=completion_tokens,
                        cached_tokens=cached_tokens,
                        extra={
                            "thoughts_tokens": thoughts_tokens,
                            "tool_tokens": tool_tokens,
                        },
                    )

                # Use thoughts as message when content is empty
                display_message = content if content else (reasoning_content or "")

                steps.append(
                    Step(
                        step_id=step_id,
                        timestamp=timestamp,
                        source="agent",
                        model_name=model_name,
                        message=display_message,
                        reasoning_content=reasoning_content if content else None,
                        tool_calls=tool_calls,
                        observation=observation,
                        metrics=metrics,
                    )
                )
                step_id += 1

        if not steps:
            return None

        # Build final metrics
        final_metrics = FinalMetrics(
            total_prompt_tokens=total_input_tokens,
            total_completion_tokens=total_output_tokens,
            total_cached_tokens=total_cached_tokens,
            total_steps=len(steps),
        )

        # Determine model name from first agent step
        default_model_name: str | None = None
        for step in steps:
            if step.source == "agent" and step.model_name:
                default_model_name = step.model_name
                break

        # Build trajectory
        trajectory = Trajectory(
            schema_version="ATIF-v1.6",
            session_id=session_id,
            agent=Agent(
                name="gemini-cli",
                version=self.version() or "unknown",
                model_name=default_model_name,
            ),
            steps=steps,
            final_metrics=final_metrics,
        )

        return trajectory

    def populate_context_post_run(self, context: AgentContext) -> None:
        gemini_path = self.logs_dir / "gemini-cli.trajectory.json"

        if not gemini_path.exists():
            return

        try:
            gemini_trajectory = json.loads(gemini_path.read_text())
        except Exception as e:
            print(f"Error loading Gemini trajectory: {e}")
            return

        # Calculate token counts for context
        n_input_tokens = 0
        n_output_tokens = 0
        n_cache_tokens = 0
        for message in gemini_trajectory.get("messages", []):
            if message.get("type") == "gemini":
                tokens = message.get("tokens", {})
                n_input_tokens += tokens.get("input", 0)
                n_output_tokens += (
                    tokens.get("output", 0)
                    + tokens.get("tool", 0)
                    + tokens.get("thoughts", 0)
                )
                n_cache_tokens += tokens.get("cached", 0)

        context.n_input_tokens = n_input_tokens
        context.n_output_tokens = n_output_tokens
        context.n_cache_tokens = n_cache_tokens

        # Convert to ATIF format and save as trajectory.json
        try:
            atif_trajectory = self._convert_gemini_to_atif(gemini_trajectory)

            if atif_trajectory:
                # Save ATIF format as trajectory.json
                atif_path = self.logs_dir / "trajectory.json"
                with open(atif_path, "w") as f:
                    json.dump(atif_trajectory.to_json_dict(), f, indent=2)
        except Exception as e:
            print(f"Error converting Gemini trajectory to ATIF: {e}")

    def create_cleanup_commands(self) -> list[ExecInput]:
        return [
            ExecInput(
                command=(
                    "find ~/.gemini/tmp -type f -name 'session-*.json' 2>/dev/null | "
                    "head -n 1 | xargs -r -I{} cp {} /logs/agent/gemini-cli.trajectory.json"
                ),
            ),
        ]

    def _build_register_skills_command(self) -> str | None:
        """Return a shell command that copies skills to Gemini CLI's skills directory."""
        if not self.skills_dir:
            return None
        return (
            f"mkdir -p ~/.gemini/skills && "
            f"cp -r {shlex.quote(self.skills_dir)}/* "
            f"~/.gemini/skills/ 2>/dev/null || true"
        )

    def _build_register_mcp_servers_command(self) -> str | None:
        """Return a shell command that writes MCP config to ~/.gemini/settings.json."""
        if not self.mcp_servers:
            return None
        servers: dict[str, dict[str, Any]] = {}
        for server in self.mcp_servers:
            if server.transport == "stdio":
                servers[server.name] = {"command": server.command, "args": server.args}
            elif server.transport == "streamable-http":
                servers[server.name] = {"httpUrl": server.url}
            else:  # sse
                servers[server.name] = {"url": server.url}
        config = json.dumps({"mcpServers": servers}, indent=2)
        escaped = shlex.quote(config)
        return f"mkdir -p ~/.gemini && echo {escaped} > ~/.gemini/settings.json"

    def create_run_agent_commands(self, instruction: str) -> list[ExecInput]:
        escaped_instruction = shlex.quote(instruction)

        if not self.model_name or "/" not in self.model_name:
            raise ValueError("Model name must be in the format provider/model_name")

        model = self.model_name.split("/")[-1]

        env = {}

        auth_vars = [
            "GEMINI_API_KEY",
            "GOOGLE_APPLICATION_CREDENTIALS",
            "GOOGLE_CLOUD_PROJECT",
            "GOOGLE_CLOUD_LOCATION",
            "GOOGLE_GENAI_USE_VERTEXAI",
            "GOOGLE_API_KEY",
        ]
        for var in auth_vars:
            if var in os.environ:
                env[var] = os.environ[var]

        commands: list[ExecInput] = []

        skills_command = self._build_register_skills_command()
        if skills_command:
            commands.append(ExecInput(command=skills_command, env=env))

        mcp_command = self._build_register_mcp_servers_command()
        if mcp_command:
            commands.append(ExecInput(command=mcp_command, env=env))

        cli_flags = self.build_cli_flags()
        extra_flags = (cli_flags + " ") if cli_flags else ""

        commands.append(
            ExecInput(
                command=(
                    ". ~/.nvm/nvm.sh; "
                    f"gemini --yolo {extra_flags}--model={model} --prompt={escaped_instruction} "
                    f"2>&1 </dev/null | stdbuf -oL tee /logs/agent/gemini-cli.txt"
                ),
                env=env,
            ),
        )

        return commands
