#!/usr/bin/env python3
"""
Workload optimization tools for CAST.AI External MCP Server.

Provides tools for workload autoscaling management, analysis, and optimization.
"""

import json
from datetime import datetime, timezone

from fastmcp import FastMCP

from ..cache import resolve_cluster_id
from ..client import make_paginated_request, make_request


def register_workload_tools(mcp: FastMCP):
    """Register all workload optimization related MCP tools."""

    @mcp.tool()
    async def get_workloads(
        cluster_id_or_name: str,
        page_limit: int | None = None,
        auto_paginate: bool = False,
        max_pages: int | None = None,
        namespaces: str | None = None,
        kinds: str | None = None,
        workload_ids: str | None = None,
        workload_names: str | None = None,
        scaling_policy_names: str | None = None,
        management_options: str | None = None,
        configured_by: str | None = None,
        search_query: str | None = None,
        sort_field: str | None = None,
        sort_order: str | None = None,
        recommendation_status: str | None = None,
        recommendation_is_low_confidence: bool | None = None,
        workload_has_error: bool | None = None,
        include_recommendations: bool = True,
        include_containers: bool = True,
        include_labels: bool = False,
        include_annotations: bool = False,
    ) -> str:
        """
        Get all workloads managed by workload autoscaling.

        Returns list of workloads with their autoscaling configuration
        and recommendations. Supports pagination for large result sets.

        Args:
            cluster_id_or_name: The cluster ID (UUID) or cluster name
            page_limit: Number of items per page (optional)
            auto_paginate: If True, automatically fetch all pages (default: False)
            max_pages: Maximum number of pages to fetch when auto_paginate is True (optional)
            namespaces: Comma-separated list of namespaces to filter by (optional)
            kinds: Comma-separated list of workload kinds to filter by (optional, e.g., "Deployment,StatefulSet")
            workload_ids: Comma-separated list of workload IDs (UUIDs) to filter by (optional)
            workload_names: Comma-separated list of workload names to filter by (optional)
            scaling_policy_names: Comma-separated list of scaling policy names to filter by (optional)
            management_options: Comma-separated list of management options to filter by (optional).
                Available values: UNDEFINED, READ_ONLY, MANAGED
            configured_by: Comma-separated list of configuration sources to filter by (optional)
            search_query: Free-text search query to filter workloads (optional)
            sort_field: Field name to sort results by (optional)
            sort_order: Sort order - "ASC" or "DESC" (optional, default: ASC)
            recommendation_status: Filter by recommendation status (optional).
                Available values: STATUS_UNKNOWN, STATUS_WAITING, STATUS_APPLIED, STATUS_STOPPED
            recommendation_is_low_confidence: If True, filter for workloads with low confidence recommendations (optional)
            workload_has_error: If True, filter for workloads that have errors (optional)
            include_recommendations: If True, include recommendation data (default: True).
                Set to False to reduce response size when only listing workloads.
            include_containers: If True, include container details (default: True).
                Set to False to reduce response size when only listing workloads.
            include_labels: If True, include workload labels (default: False).
                Set to True when you need to filter or analyze workloads by labels.
            include_annotations: If True, include workload annotations (default: False).
                Set to True when you need to see workload annotations.

        Returns:
            JSON string with workload autoscaling data
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)

            # Build params - API expects arrays for multi-value params
            params: dict = {}
            if namespaces:
                params["namespaces"] = [n.strip() for n in namespaces.split(",")]
            if kinds:
                params["kinds"] = [k.strip() for k in kinds.split(",")]
            if workload_ids:
                params["workloadIds"] = [w.strip() for w in workload_ids.split(",")]
            if workload_names:
                params["workloadNames"] = [n.strip() for n in workload_names.split(",")]
            if scaling_policy_names:
                params["scalingPolicyNames"] = [p.strip() for p in scaling_policy_names.split(",")]
            if management_options:
                params["managementOptions"] = [m.strip().upper() for m in management_options.split(",")]
            if configured_by:
                params["configuredBy"] = [c.strip() for c in configured_by.split(",")]
            if search_query:
                params["searchQuery"] = search_query
            if sort_field:
                params["sort.field"] = sort_field
            if sort_order:
                params["sort.order"] = sort_order.upper()
            if recommendation_status:
                params["recommendationStatusType"] = recommendation_status.upper()
            if recommendation_is_low_confidence is not None:
                params["recommendationIsLowConfidence"] = recommendation_is_low_confidence
            if workload_has_error is not None:
                params["workloadHasError"] = workload_has_error

            result = await make_paginated_request(
                "GET",
                f"/v1/workload-autoscaling/clusters/{cluster_id}/workloads",
                params=params if params else None,
                page_limit=page_limit,
                auto_paginate=auto_paginate,
                max_pages=max_pages,
            )

            # Strip large objects if not requested to reduce response size
            if "items" in result:
                for item in result["items"]:
                    if not include_recommendations and "recommendation" in item:
                        del item["recommendation"]
                    if not include_containers and "containers" in item:
                        del item["containers"]
                    if not include_labels and "labels" in item:
                        del item["labels"]
                    if not include_annotations and "annotations" in item:
                        del item["annotations"]

            return json.dumps(result, indent=2)
        except Exception as e:
            return json.dumps({"error": str(e)})

    @mcp.tool()
    async def get_workload_summary(cluster_id_or_name: str) -> str:
        """
        Get summary of workload autoscaling for a cluster.

        Provides overview of autoscaling status and metrics.

        Args:
            cluster_id_or_name: The cluster ID (UUID) or cluster name

        Returns:
            JSON string with workload autoscaling summary
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)
            result = await make_request(
                "GET", f"/v1/workload-autoscaling/clusters/{cluster_id}/workloads-summary"
            )
            return json.dumps(result, indent=2)
        except Exception as e:
            return json.dumps({"error": str(e)})

    @mcp.tool()
    async def analyze_workload_opportunities(
        cluster_id_or_name: str,
        top_n: int = 20,
        min_savings_percent: int | None = None,
        namespaces: str | None = None,
    ) -> str:
        """
        Analyze and rank the biggest workload autoscaling optimization opportunities.

        This tool fetches workload data and analyzes it to identify the workloads with
        the highest potential for resource optimization (CPU and memory savings).

        Args:
            cluster_id_or_name: The cluster ID (UUID) or cluster name
            top_n: Number of top opportunities to return (default: 20)
            min_savings_percent: Minimum savings percentage to include (optional)
            namespaces: Comma-separated list of namespaces to analyze (optional)

        Returns:
            JSON string with ranked optimization opportunities including:
            - Workload details (name, namespace, kind)
            - Current vs recommended resources (CPU and memory)
            - Potential savings (absolute and percentage)
            - Optimization status and recommendations
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)

            # Build params for filtering
            params = {}
            if namespaces:
                params["namespaces"] = namespaces

            # Fetch all workloads (with auto-pagination to get complete data)
            result = await make_paginated_request(
                "GET",
                f"/v1/workload-autoscaling/clusters/{cluster_id}/workloads",
                params=params,
                auto_paginate=True,
            )

            workloads = result.get("workloads", [])

            if not workloads:
                return json.dumps({"message": "No workloads found", "opportunities": []})

            # Analyze each workload for optimization opportunities
            opportunities = []

            for workload in workloads:
                # Extract key fields
                name = workload.get("name", "unknown")
                namespace = workload.get("namespace", "unknown")
                kind = workload.get("kind", "unknown")
                replicas = workload.get("replicas", 0)

                # Get current and recommended resources
                containers = workload.get("containers", [])

                total_current_cpu = 0
                total_recommended_cpu = 0
                total_current_memory = 0
                total_recommended_memory = 0
                has_recommendations = False

                for container in containers:
                    # Current requests
                    resources = container.get("resources", {})
                    current_requests = resources.get("requests")

                    if current_requests is None:
                        original_resources = container.get("originalResources", {})
                        current_requests = original_resources.get("requests")

                    if current_requests and current_requests is not None:
                        current_cpu_cores = current_requests.get("cpuCores", 0)
                        current_memory_gib = current_requests.get("memoryGib", 0)
                    else:
                        current_cpu_cores = 0
                        current_memory_gib = 0

                    # Recommendations
                    recommendation = container.get("recommendation")

                    if recommendation and recommendation is not None:
                        rec_requests = recommendation.get("requests", {})
                        if rec_requests:
                            rec_cpu_cores = rec_requests.get("cpuCores", 0)
                            rec_memory_gib = rec_requests.get("memoryGib", 0)
                            has_recommendations = True
                        else:
                            rec_cpu_cores = 0
                            rec_memory_gib = 0
                    else:
                        rec_cpu_cores = 0
                        rec_memory_gib = 0

                    # Convert to millicores and MiB
                    total_current_cpu += current_cpu_cores * 1000
                    total_recommended_cpu += rec_cpu_cores * 1000
                    total_current_memory += current_memory_gib * 1024
                    total_recommended_memory += rec_memory_gib * 1024

                # Calculate savings (accounting for replicas)
                cpu_savings = (total_current_cpu - total_recommended_cpu) * replicas
                memory_savings = (total_current_memory - total_recommended_memory) * replicas

                # Calculate percentages
                cpu_savings_percent = (cpu_savings / total_current_cpu * 100) if total_current_cpu > 0 else 0
                memory_savings_percent = (
                    (memory_savings / total_current_memory * 100) if total_current_memory > 0 else 0
                )

                # Determine opportunity type
                has_savings_opportunity = False
                opportunity_type = "reduction"

                if total_current_cpu == 0 and total_recommended_cpu > 0:
                    has_savings_opportunity = True
                    opportunity_type = "set_requests"
                    cpu_savings = -total_recommended_cpu * replicas
                elif cpu_savings > 0:
                    has_savings_opportunity = True
                    opportunity_type = "reduction"

                if total_current_memory == 0 and total_recommended_memory > 0:
                    has_savings_opportunity = True
                    if opportunity_type != "reduction":
                        opportunity_type = "set_requests"
                    memory_savings = -total_recommended_memory * replicas
                elif memory_savings > 0:
                    has_savings_opportunity = True
                    if opportunity_type == "set_requests":
                        opportunity_type = "mixed"

                # Only include workloads with savings or recommendations
                if has_savings_opportunity or cpu_savings > 0 or memory_savings > 0:
                    # Calculate savings score for ranking
                    savings_score = (cpu_savings / 1000) + (memory_savings / 1024 * 0.5)

                    opportunity = {
                        "workload": {
                            "name": name,
                            "namespace": namespace,
                            "kind": kind,
                            "replicas": replicas,
                        },
                        "current": {
                            "cpu_millicores": round(total_current_cpu, 2),
                            "memory_mib": round(total_current_memory, 2),
                        },
                        "recommended": {
                            "cpu_millicores": round(total_recommended_cpu, 2),
                            "memory_mib": round(total_recommended_memory, 2),
                        },
                        "savings_potential": {
                            "cpu_millicores": round(cpu_savings, 2),
                            "cpu_percent": round(cpu_savings_percent, 1),
                            "memory_mib": round(memory_savings, 2),
                            "memory_percent": round(memory_savings_percent, 1),
                            "score": round(savings_score, 2),
                            "type": opportunity_type,
                        },
                        "optimization_status": workload.get("optimizationStatus", "unknown"),
                    }

                    # Apply min_savings_percent filter if specified
                    if (
                        min_savings_percent is None
                        or max(cpu_savings_percent, memory_savings_percent) >= min_savings_percent
                    ):
                        opportunities.append(opportunity)

            # Sort by savings score (highest first)
            opportunities.sort(key=lambda x: x["savings_potential"]["score"], reverse=True)

            # Limit to top N
            top_opportunities = opportunities[:top_n]

            # Calculate totals
            total_cpu_savings = sum(o["savings_potential"]["cpu_millicores"] for o in top_opportunities)
            total_memory_savings = sum(o["savings_potential"]["memory_mib"] for o in top_opportunities)

            response = {
                "cluster_id": cluster_id,
                "total_workloads_analyzed": len(workloads),
                "opportunities_found": len(opportunities),
                "top_opportunities_shown": len(top_opportunities),
                "total_savings_potential": {
                    "cpu_cores": round(total_cpu_savings / 1000, 2),
                    "memory_gib": round(total_memory_savings / 1024, 2),
                },
                "opportunities": top_opportunities,
            }

            return json.dumps(response, indent=2)

        except Exception as e:
            return json.dumps({"error": str(e)})

    @mcp.tool()
    async def get_workload(
        cluster_id_or_name: str,
        workload_id: str,
        include_metrics: bool = False,
        include_costs: bool = False,
        include_labels: bool = False,
        include_annotations: bool = False,
        from_time: str | None = None,
        to_time: str | None = None,
    ) -> str:
        """
        Get detailed information about a specific workload by its ID.

        Returns comprehensive workload data including configuration, resource requests/limits,
        autoscaling settings, recommendations, and optimization status.

        Use this when you have a workload ID and need full details about that specific workload.
        To find workload IDs, use get_workloads() first.

        Args:
            cluster_id_or_name: The cluster ID (UUID) or cluster name
            workload_id: The unique identifier of the workload (UUID)
            include_metrics: Include workload metrics in response (default: False)
            include_costs: Include cost data in response (default: False)
            include_labels: Include workload labels in response (default: False)
            include_annotations: Include workload annotations in response (default: False)
            from_time: Start time for metrics/costs in RFC3339 format (e.g., "2025-01-01T00:00:00Z")
            to_time: End time for metrics/costs in RFC3339 format (e.g., "2025-01-31T23:59:59Z")

        Returns:
            JSON string with detailed workload information
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)

            # Get current time for validation
            now = datetime.now(timezone.utc)

            # Validate and clamp to_time to current time if in the future
            validated_to_time = to_time
            if to_time:
                try:
                    to_time_str = to_time.replace("Z", "+00:00")
                    parsed_to_time = datetime.fromisoformat(to_time_str)
                    if parsed_to_time > now:
                        validated_to_time = now.strftime("%Y-%m-%dT%H:%M:%SZ")
                except ValueError:
                    pass

            # Validate from_time is not in the future
            validated_from_time = from_time
            if from_time:
                try:
                    from_time_str = from_time.replace("Z", "+00:00")
                    parsed_from_time = datetime.fromisoformat(from_time_str)
                    if parsed_from_time > now:
                        validated_from_time = now.strftime("%Y-%m-%dT%H:%M:%SZ")
                except ValueError:
                    pass

            params: dict = {}
            if include_metrics:
                params["includeMetrics"] = "true"
            if include_costs:
                params["includeCosts"] = "true"
            if validated_from_time:
                params["fromTime"] = validated_from_time
            if validated_to_time:
                params["toTime"] = validated_to_time

            result = await make_request(
                "GET",
                f"/v1/workload-autoscaling/clusters/{cluster_id}/workloads/{workload_id}",
                params=params if params else None,
            )

            # Strip labels/annotations if not requested
            if not include_labels and "labels" in result:
                del result["labels"]
            if not include_annotations and "annotations" in result:
                del result["annotations"]

            return json.dumps(result, indent=2)
        except Exception as e:
            return json.dumps({"error": str(e)})

    @mcp.tool()
    async def get_workload_events(
        cluster_id_or_name: str,
        from_date: str | None = None,
        to_date: str | None = None,
        workload_id: str | None = None,
        workload_name: str | None = None,
        workload_namespace: str | None = None,
        workload_kind: str | None = None,
        event_types: str | None = None,
        page_limit: int | None = None,
        auto_paginate: bool = False,
        max_pages: int | None = None,
        include_event_data: bool = False,
    ) -> str:
        """
        Get workload autoscaling events for a cluster.

        Returns a list of events related to workload autoscaling operations such as
        OOM kills, surges, configuration changes, recommendation updates, and more.

        Args:
            cluster_id_or_name: The cluster ID (UUID) or cluster name
            from_date: Filter events from this date (RFC3339 format, e.g., "2025-01-01T00:00:00Z")
            to_date: Filter events until this date (RFC3339 format, e.g., "2025-01-31T23:59:59Z")
            workload_id: Filter by specific workload ID (optional)
            workload_name: Filter by workload name (optional)
            workload_namespace: Filter by namespace (optional)
            workload_kind: Filter by workload kind, e.g., "Deployment", "StatefulSet" (optional)
            event_types: Comma-separated list of event types to filter (optional).
                Available types: SURGE, OOM_KILL, CONFIGURATION_CHANGEDV2,
                RECOMMENDED_POD_COUNT_CHANGED, RECOMMENDED_REQUESTS_CHANGED,
                SCALING_POLICY_CREATED, SCALING_POLICY_DELETED, SCALING_POLICY_UPDATED,
                SCALING_POLICY_ASSIGNED, FAILED_HELM_TEST_HOOK, SCALING_POLICY_ORDER_UPDATED,
                MEMORY_PRESSURE_EVICTION, SYSTEM_OVERRIDE_TRIGGERED, SYSTEM_OVERRIDE_RESET,
                STARTUP_FAILURE
            page_limit: Number of items per page (optional)
            auto_paginate: If True, automatically fetch all pages (default: False)
            max_pages: Maximum number of pages to fetch when auto_paginate is True (optional)
            include_event_data: If True, include the full event object with detailed data (default: False)

        Returns:
            JSON string with workload events data
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)

            # Get current time for validation
            now = datetime.now(timezone.utc)

            # Validate and clamp to_date to current time if in the future
            validated_to_date = to_date
            if to_date:
                try:
                    to_date_str = to_date.replace("Z", "+00:00")
                    parsed_to_date = datetime.fromisoformat(to_date_str)
                    if parsed_to_date > now:
                        validated_to_date = now.strftime("%Y-%m-%dT%H:%M:%SZ")
                except ValueError:
                    pass

            # Build params
            params: dict = {}
            if from_date:
                params["fromDate"] = from_date
            if validated_to_date:
                params["toDate"] = validated_to_date
            if workload_id:
                params["workloadId"] = workload_id
            if workload_name:
                params["workloadName"] = workload_name
            if workload_namespace:
                params["workloadNamespace"] = workload_namespace
            if workload_kind:
                params["workloadKind"] = workload_kind
            if event_types:
                # Convert comma-separated types to list with EVENT_TYPE_ prefix
                types_list = [
                    (
                        f"EVENT_TYPE_{t.strip().upper()}"
                        if not t.strip().upper().startswith("EVENT_TYPE_")
                        else t.strip().upper()
                    )
                    for t in event_types.split(",")
                ]
                params["type"] = types_list

            result = await make_paginated_request(
                "GET",
                f"/v1/workload-autoscaling/clusters/{cluster_id}/workload-events",
                params=params,
                page_limit=page_limit,
                auto_paginate=auto_paginate,
                max_pages=max_pages,
            )

            # Strip verbose event data if not requested
            if not include_event_data and "items" in result:
                for item in result["items"]:
                    if "event" in item:
                        del item["event"]

            return json.dumps(result, indent=2)
        except Exception as e:
            return json.dumps({"error": str(e)})

    @mcp.tool()
    async def list_workload_scaling_policies(
        cluster_id_or_name: str,
        page_limit: int | None = None,
        auto_paginate: bool = False,
        max_pages: int | None = None,
    ) -> str:
        """
        List all workload scaling policies for a cluster.

        Returns a list of scaling policies that define how workloads should be
        autoscaled, including resource requests, limits, and scaling behavior.

        Args:
            cluster_id_or_name: The cluster ID (UUID) or cluster name
            page_limit: Number of items per page (optional)
            auto_paginate: If True, automatically fetch all pages (default: False)
            max_pages: Maximum number of pages to fetch when auto_paginate is True (optional)

        Returns:
            JSON string with list of scaling policies
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)
            result = await make_paginated_request(
                "GET",
                f"/v1/workload-autoscaling/clusters/{cluster_id}/policies",
                page_limit=page_limit,
                auto_paginate=auto_paginate,
                max_pages=max_pages,
            )
            return json.dumps(result, indent=2)
        except Exception as e:
            return json.dumps({"error": str(e)})

    @mcp.tool()
    async def get_workload_scaling_policy(
        cluster_id_or_name: str,
        policy_id: str,
    ) -> str:
        """
        Get detailed information about a specific workload scaling policy.

        Returns comprehensive policy data including resource configuration,
        scaling rules, workload selectors, and policy metadata.

        Args:
            cluster_id_or_name: The cluster ID (UUID) or cluster name
            policy_id: The unique identifier of the scaling policy (UUID)

        Returns:
            JSON string with detailed scaling policy information
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)
            result = await make_request(
                "GET",
                f"/v1/workload-autoscaling/clusters/{cluster_id}/policies/{policy_id}",
            )
            return json.dumps(result, indent=2)
        except Exception as e:
            return json.dumps({"error": str(e)})

    @mcp.tool()
    async def list_workload_fields(cluster_id_or_name: str) -> str:
        """
        List all unique field names found in workload responses.

        Useful for discovering what fields are available in the API response.

        Args:
            cluster_id_or_name: The cluster ID (UUID) or cluster name

        Returns:
            JSON string with all field names found at various levels
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)

            # Fetch first page of workloads
            result = await make_paginated_request(
                "GET",
                f"/v1/workload-autoscaling/clusters/{cluster_id}/workloads",
                page_limit=10,
            )

            workloads = result.get("workloads", [])

            if not workloads:
                return json.dumps({"error": "No workloads found"})

            # Collect all unique field names
            workload_fields = set()
            container_fields = set()

            for w in workloads:
                workload_fields.update(w.keys())
                for container in w.get("containers", []):
                    container_fields.update(container.keys())

            return json.dumps(
                {
                    "workload_level_fields": sorted(workload_fields),
                    "container_level_fields": sorted(container_fields),
                    "sample_workload_names": [w.get("name") for w in workloads[:5]],
                },
                indent=2,
            )

        except Exception as e:
            return json.dumps({"error": str(e)})

    @mcp.tool()
    async def debug_workload_structure(
        cluster_id_or_name: str,
        workload_name: str | None = None,
        show_all: bool = False,
    ) -> str:
        """
        Debug tool to inspect the actual structure of workload data.

        Use this to see what fields are available in the workload response.

        Args:
            cluster_id_or_name: The cluster ID (UUID) or cluster name
            workload_name: Optional workload name to inspect (returns first workload if not specified)
            show_all: If True, show all workloads with recommendations (default: False)

        Returns:
            JSON string with the structure of one or more workloads
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)

            # Fetch workloads
            result = await make_paginated_request(
                "GET",
                f"/v1/workload-autoscaling/clusters/{cluster_id}/workloads",
                page_limit=100 if show_all else 10,
                auto_paginate=show_all,
            )

            workloads = result.get("workloads", [])

            if not workloads:
                return json.dumps({"error": "No workloads found"})

            if show_all:
                # Show all workloads that have recommendations
                workloads_with_recommendations = []
                for w in workloads:
                    containers = w.get("containers", [])
                    has_rec = any(c.get("recommendation") is not None for c in containers)
                    if has_rec:
                        workloads_with_recommendations.append(
                            {
                                "name": w.get("name"),
                                "namespace": w.get("namespace"),
                                "kind": w.get("kind"),
                                "replicas": w.get("replicas"),
                                "podCount": w.get("podCount"),
                                "containers": [
                                    {
                                        "name": c.get("name"),
                                        "has_recommendation": c.get("recommendation") is not None,
                                        "recommendation": c.get("recommendation"),
                                        "requests": c.get("requests"),
                                    }
                                    for c in containers
                                ],
                            }
                        )

                return json.dumps(
                    {
                        "message": f"Found {len(workloads_with_recommendations)} workloads with recommendations",
                        "total_workloads": len(workloads),
                        "workloads": workloads_with_recommendations,
                    },
                    indent=2,
                )

            # Find specific workload or return first one
            target_workload = workloads[0]
            if workload_name:
                for w in workloads:
                    if w.get("name") == workload_name:
                        target_workload = w
                        break

            return json.dumps(
                {"message": "Workload structure for debugging", "workload": target_workload},
                indent=2,
            )

        except Exception as e:
            return json.dumps({"error": str(e)})