"""
Utility functions for the Anthropic API.
"""

import base64
import io
import os
import sys
from dataclasses import dataclass, fields, replace
from enum import Enum
from typing import List, Optional, Union, cast

import boto3
from anthropic.types.beta import (
    BetaCacheControlEphemeralParam,
    BetaImageBlockParam,
    BetaMessage,
    BetaMessageParam,
    BetaTextBlockParam,
    BetaToolResultBlockParam,
)
from PIL import Image

COMPUTER_USE_BETA_FLAG = "computer-use-2025-11-24"
PROMPT_CACHING_BETA_FLAG = "prompt-caching-2024-07-31"
COMPUTER_USE_TYPE = "computer_20251124"

# computer_20250124 computer-use-2025-01-24
# computer_20251124 computer-use-2025-11-24


def resize_and_encode_png(
    screenshot_bytes: bytes, size: tuple[int, int] = (1280, 720)
) -> str:
    """Resize PNG bytes to *size* and return a base64-encoded PNG string."""
    if not screenshot_bytes:
        return ""
    img = Image.open(io.BytesIO(screenshot_bytes))
    resized = img.resize(size, Image.Resampling.LANCZOS)
    buf = io.BytesIO()
    resized.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")


@dataclass(frozen=True)  # kw_only=True,
class ToolResult:
    """Represents the result of a tool execution."""

    output: Optional[str] = None
    error: Optional[str] = None
    base64_data: Optional[str] = None
    media_type: Optional[str] = None
    system: Optional[str] = None

    def __bool__(self):
        return any(getattr(self, field.name) for field in fields(self))

    def __add__(self, other: "ToolResult"):
        def combine_fields(
            field: Optional[str], other_field: Optional[str], concatenate: bool = True
        ):
            if field and other_field:
                if concatenate:
                    return field + other_field
                raise ValueError("Cannot combine tool results")
            return field or other_field

        return ToolResult(
            output=combine_fields(self.output, other.output),
            error=combine_fields(self.error, other.error),
            base64_data=combine_fields(self.base64_data, other.base64_data, False),
            media_type=combine_fields(self.media_type, other.media_type, False),
            system=combine_fields(self.system, other.system),
        )

    def replace(self, **kwargs):
        """Returns a new ToolResult with the given fields replaced."""
        return replace(self, **kwargs)


class APIProvider(Enum):
    ANTHROPIC = "anthropic"
    BEDROCK = "bedrock"
    VERTEX = "vertex"


PROVIDER_TO_DEFAULT_MODEL_NAME: dict[(APIProvider, str), str] = {
    (APIProvider.ANTHROPIC, "claude-sonnet-4-6"): "claude-sonnet-4-6",
    (APIProvider.BEDROCK, "claude-sonnet-4-6"): "global.anthropic.claude-sonnet-4-6",
    (APIProvider.ANTHROPIC, "claude-opus-4-6"): "claude-opus-4-6",
    (APIProvider.BEDROCK, "claude-opus-4-6"): "global.anthropic.claude-opus-4-6-v1",
    (APIProvider.ANTHROPIC, "claude-opus-4-7"): "claude-opus-4-7",
    (APIProvider.BEDROCK, "claude-opus-4-7"): "global.anthropic.claude-opus-4-7",
    (APIProvider.ANTHROPIC, "claude-haiku-4-5"): "claude-haiku-4-5-20251001",
    (
        APIProvider.BEDROCK,
        "claude-haiku-4-5",
    ): "global.anthropic.claude-haiku-4-5-20251001-v1:0",
}


def get_model_name(provider: APIProvider, model_name: str) -> str:
    """
    Get the actual model name to use for API calls.

    Simply returns the model name as-is for direct API usage.
    """
    # Look up in the mapping table
    key = (provider, model_name)
    if key in PROVIDER_TO_DEFAULT_MODEL_NAME:
        return PROVIDER_TO_DEFAULT_MODEL_NAME[key]

    # If not found and using Anthropic direct API, return as-is (may work for new models)
    if provider == APIProvider.ANTHROPIC:
        return model_name

    # For Bedrock/Vertex, show available models if not found
    available_models = [
        k[1] for k in PROVIDER_TO_DEFAULT_MODEL_NAME.keys() if k[0] == provider
    ]
    raise ValueError(
        f"❌ Model '{model_name}' is not configured for {provider.value}.\n"
        f"📋 Available models for {provider.value}: {available_models}"
    )


def _make_api_tool_result(
    result: ToolResult, tool_use_id: str
) -> BetaToolResultBlockParam:
    """Convert an agent ToolResult to an API ToolResultBlockParam."""
    tool_result_content: Union[
        List[Union[BetaTextBlockParam, BetaImageBlockParam]], str
    ] = []
    is_error = False

    if not result or (result.get("error") is not None and result.get("error") != ""):
        is_error = True
        error_message = (
            str(result.get("error", "Unknown error occurred"))
            if result
            else "No result received"
        )
        tool_result_content = [
            {
                "type": "text",
                "text": _maybe_prepend_system_tool_result(result, error_message),
            }
        ]

    else:
        if result.get("output"):
            tool_result_content.append(
                {
                    "type": "text",
                    "text": _maybe_prepend_system_tool_result(
                        result, str(result.get("output", "") if result else "")
                    ),
                }
            )

        if result.get("base64_data") and result.get("media_type"):
            media_type = result["media_type"]
            if media_type == "application/pdf":
                tool_result_content.append(
                    {
                        "type": "document",
                        "source": {
                            "type": "base64",
                            "media_type": "application/pdf",
                            "data": result["base64_data"],
                        },
                    }
                )
            elif media_type.startswith("image/"):
                tool_result_content.append(
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": media_type,
                            "data": result["base64_data"],
                        },
                    }
                )

        if not tool_result_content:
            tool_result_content.append(
                {"type": "text", "text": "Action completed successfully"}
            )

    return {
        "type": "tool_result",
        "content": tool_result_content,
        "tool_use_id": tool_use_id,
        "is_error": is_error,
    }


def _maybe_prepend_system_tool_result(result: ToolResult, result_text: str):
    if not result:
        return result_text

    if result.get("system", False):
        result_text = f"<system>{result.get('system', '')}</system>\n{result_text}"
    return result_text


def _inject_prompt_caching(
    messages: list[BetaMessageParam],
):
    """
    Set cache breakpoints for the 3 most recent turns
    one cache breakpoint is left for tools/system prompt, to be shared across sessions
    """

    breakpoints_remaining = 2  # Use full budget for recent messages
    messages_processed = 0

    for message in reversed(messages):
        if message["role"] == "user" and isinstance(
            content := message["content"], list
        ):
            messages_processed += 1
            # Check if this message would fit within the remaining budget
            if breakpoints_remaining >= len(content):
                # We have enough budget, spend it and add cache_control
                breakpoints_remaining -= len(content)
                # Use type ignore to bypass TypedDict check until SDK types are updated
                content[-1]["cache_control"] = BetaCacheControlEphemeralParam(  # type: ignore
                    {"type": "ephemeral"}
                )
            else:
                # Check if this is the first message (contains image + text with task description)
                is_first_message = messages_processed == len(
                    [msg for msg in messages if msg["role"] == "user"]
                )

                if not is_first_message:
                    # Not enough budget, remove any existing cache_control from this message
                    content[-1].pop("cache_control", None)
                # Continue to clean up older messages that might have cache_control from previous turns


def _maybe_filter_to_n_most_recent_images(
    messages: list[BetaMessageParam],
    images_to_keep: int,
    min_removal_threshold: int,
):
    """
    With the assumption that images are screenshots that are of diminishing value as
    the conversation progresses, remove all but the final `images_to_keep` tool_result
    images in place, with a chunk of min_removal_threshold to reduce the amount we
    break the implicit prompt cache.
    """
    if images_to_keep is None:
        return messages

    tool_result_blocks = cast(
        list[BetaToolResultBlockParam],
        [
            item
            for message in messages
            for item in (
                message["content"] if isinstance(message["content"], list) else []
            )
            if isinstance(item, dict) and item.get("type") == "tool_result"
        ],
    )

    total_images = sum(
        1
        for tool_result in tool_result_blocks
        for content in tool_result.get("content", [])
        if isinstance(content, dict) and content.get("type") == "image"
    )

    images_to_remove = total_images - images_to_keep
    # for better cache behavior, we want to remove in chunks
    images_to_remove -= images_to_remove % min_removal_threshold

    for tool_result in tool_result_blocks:
        if isinstance(tool_result.get("content"), list):
            new_content = []
            for content in tool_result.get("content", []):
                if isinstance(content, dict) and content.get("type") == "image":
                    if images_to_remove > 0:
                        images_to_remove -= 1
                        continue
                new_content.append(content)
            tool_result["content"] = new_content


def validate_model_support(
    model_name: str,
    provider: APIProvider = APIProvider.BEDROCK,
    api_key: str = None,
    temperature: float = None,
    top_p: float = None,
) -> bool:
    """
    Validate model support with the same API call pattern as the main agent.

    Args:
        model_name: The model name to validate
        provider: API provider (ANTHROPIC, BEDROCK, or VERTEX)
        api_key: Optional API key (only for ANTHROPIC provider), defaults to ANTHROPIC_API_KEY env var
        temperature: Optional temperature parameter for testing
        top_p: Optional top_p parameter for testing

    Returns:
        True if model is supported and API call succeeds, False otherwise
    """
    print(f"🔍 Validating model support: {model_name} (provider: {provider.value})")

    try:
        import os
        import time

        from anthropic import Anthropic, AnthropicBedrock, AnthropicVertex

        # Same client setup as main agent - choose based on provider
        if provider == APIProvider.ANTHROPIC:
            client = Anthropic(
                api_key=api_key or os.environ.get("ANTHROPIC_API_KEY"), max_retries=4
            ).with_options(default_headers={"anthropic-beta": COMPUTER_USE_BETA_FLAG})
        elif provider == APIProvider.BEDROCK:
            client = AnthropicBedrock(
                aws_access_key=os.getenv("AWS_ACCESS_KEY_ID"),
                aws_secret_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
                aws_region=os.getenv("AWS_DEFAULT_REGION"),
            )
        elif provider == APIProvider.VERTEX:
            client = AnthropicVertex()

        # Same message format as main agent - always use structured format with cache_control
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "Respond with 'OK'",
                        "cache_control": {"type": "ephemeral"},
                    }
                ],
            }
        ]

        # Same betas configuration as main agent
        betas = [COMPUTER_USE_BETA_FLAG]

        system = [{"type": "text", "text": "You are Claude. Respond with 'OK'."}]

        # Same tools configuration as main agent - use modern computer tool for all models
        tools = [
            {
                "name": "computer",
                "type": COMPUTER_USE_TYPE,
                "display_width_px": 1280,
                "display_height_px": 720,
                "display_number": 1,
            }
        ]

        # Same thinking configuration as main agent
        max_tokens = 50  # Base validation max_tokens
        extra_body = {"thinking": {"type": "adaptive"}}
        actual_max_tokens = max_tokens

        # Sampling parameters (same logic as main agent)
        sampling_params = {}
        if temperature is not None:
            sampling_params["temperature"] = temperature
        if top_p is not None:
            sampling_params["top_p"] = top_p

        # Retry logic with 5 attempts, 5 second delays
        for attempt in range(5):
            try:
                # Same API call pattern as main agent
                client.beta.messages.create(
                    max_tokens=actual_max_tokens,
                    messages=messages,
                    model=get_model_name(provider, model_name),
                    system=system,
                    tools=tools,
                    betas=betas,
                    extra_body=extra_body,
                    **sampling_params,
                )

                print(
                    f"✅ Model {model_name} validated successfully with {provider.value}"
                )
                return True
            except Exception as e:
                error_msg = str(e)
                if attempt < 4:  # Don't print error on final attempt
                    print(
                        f'🔄 Validation attempt {attempt + 1}/5 failed: "{error_msg}"'
                    )
                    print("⏳ Retrying in 5 seconds...")
                    time.sleep(5)
                else:
                    print(
                        f'❌ All validation attempts failed. Final error: "{error_msg}"'
                    )

        return False

    except ValueError:
        return False
    except Exception as e:
        print(f"❌ API validation setup failed: {e}")
        return False


def _response_to_params(response: BetaMessage) -> list[dict]:
    if not response.content:
        return []
    res = []
    for block in response.content:
        dumped = {k: v for k, v in block.model_dump().items() if v is not None}
        res.append(dumped)
    if res[-1].get("type") in ("thinking", "redacted_thinking"):
        res.append({"type": "text", "text": "."})
    return res


# MARK: Validate required API keys
def _validate_parallel_api_key() -> None:
    if not os.environ.get("PARALLEL_API_KEY"):
        print("ERROR: PARALLEL_API_KEY environment variable is not set.")
        sys.exit(1)


def _validate_anthropic_api_key() -> None:
    if not os.environ.get("ANTHROPIC_API_KEY"):
        print("ERROR: ANTHROPIC_API_KEY environment variable is not set.")
        sys.exit(1)


# MARK: Get Bedrock Access
def _load_aws_bedrock_access() -> dict[str, str]:
    """Assume a cross-account role and return temporary Bedrock credentials."""
    role_arn = os.environ.get("BEDROCK_ROLE_ARN")
    if not role_arn:
        print("ERROR: BEDROCK_ROLE_ARN is not set.")
        sys.exit(1)

    region = os.environ.get("AWS_BEDROCK_REGION", "us-east-1")

    try:
        sts = boto3.client("sts")
        creds = sts.assume_role(
            RoleArn=role_arn,
            RoleSessionName="pointer-bedrock-session",
        )["Credentials"]

        return {
            "aws_access_key": creds["AccessKeyId"],
            "aws_secret_key": creds["SecretAccessKey"],
            "aws_session_token": creds["SessionToken"],
            "aws_region": region,
        }
    except Exception as e:
        print(f"ERROR: Failed to assume Bedrock role {role_arn}: {e}")
        sys.exit(1)
