"""Estimate data size for a STAC query."""

import importlib.util
from typing import Any

from mcp.types import TextContent

from stac_mcp.tools import MAX_ASSET_LIST
from stac_mcp.tools.client import STACClient

try:
    ODC_STAC_AVAILABLE = (
        importlib.util.find_spec("odc.stac") is not None
    )  # pragma: no cover
except ModuleNotFoundError:  # pragma: no cover
    ODC_STAC_AVAILABLE = False


def handle_estimate_data_size(
    client: STACClient,
    arguments: dict[str, Any],
) -> list[TextContent] | dict[str, Any]:
    collections = arguments.get("collections")
    bbox = arguments.get("bbox")
    dt = arguments.get("datetime")
    query = arguments.get("query")
    aoi_geojson = arguments.get("aoi_geojson")
    limit = arguments.get("limit", 100)
    size_estimate = client.estimate_data_size(
        collections=collections,
        bbox=bbox,
        datetime=dt,
        query=query,
        aoi_geojson=aoi_geojson,
        limit=limit,
    )
    if arguments.get("output_format") == "json":
        return {"type": "data_size_estimate", "estimate": size_estimate}
    result_text = "**Data Size Estimation**\n\n"
    result_text += f"Items analyzed: {size_estimate['item_count']}\n"
    result_text += (
        f"Estimated size: {size_estimate['estimated_size_mb']:.2f} MB "
        f"({size_estimate['estimated_size_gb']:.4f} GB)\n"
    )
    result_text += f"Raw bytes: {size_estimate['estimated_size_bytes']:,}\n\n"
    result_text += "**Query Parameters:**\n"
    result_text += "Collections: "
    collections_list = (
        ", ".join(size_estimate["collections"])
        if size_estimate["collections"]
        else "All"
    )
    result_text += f"{collections_list}\n"
    if size_estimate["bbox_used"]:
        b = size_estimate["bbox_used"]
        result_text += (
            f"Bounding box: [{b[0]:.4f}, {b[1]:.4f}, {b[2]:.4f}, {b[3]:.4f}]\n"
        )
    if size_estimate["temporal_extent"]:
        result_text += f"Time range: {size_estimate['temporal_extent']}\n"
    if size_estimate["clipped_to_aoi"]:
        result_text += "Clipped to AOI: Yes (minimized to smallest area)\n"
    if "data_variables" in size_estimate:
        result_text += "\n**Data Variables:**\n"
        for var_info in size_estimate["data_variables"]:
            result_text += (
                f"  - {var_info['variable']}: {var_info['size_mb']} MB, "
                f"shape {var_info['shape']}, dtype {var_info['dtype']}\n"
            )
    if size_estimate.get("spatial_dims"):
        spatial = size_estimate["spatial_dims"]
        result_text += "\n**Spatial Dimensions:**\n"
        result_text += f"  X (longitude): {spatial.get('x', 0)} pixels\n"
        result_text += f"  Y (latitude): {spatial.get('y', 0)} pixels\n"
    if "assets_analyzed" in size_estimate:
        result_text += "\n**Assets Analyzed (fallback estimation):**\n"
        for asset_info in size_estimate["assets_analyzed"][:MAX_ASSET_LIST]:
            result_text += (
                f"  - {asset_info['asset']}: {asset_info['estimated_size_mb']} MB "
                f"({asset_info['media_type']})\n"
            )
        remaining = len(size_estimate["assets_analyzed"]) - MAX_ASSET_LIST
        if remaining > 0:
            result_text += f"  ... and {remaining} more assets\n"
    result_text += f"\n{size_estimate['message']}\n"
    return [TextContent(type="text", text=result_text)]
