#!/usr/bin/env python3
"""
Autoscaler policy update tool for CAST.AI External MCP Server.

Allows enabling/disabling the cluster autoscaler and configuring its core policies
(unschedulable pods, node downscaler, cluster limits, spot instances).
"""

import json
from typing import Any

from fastmcp import FastMCP

from ..cache import resolve_cluster_id
from ..client import make_request
from ..logger import logger


def assess_autoscaler_risk(current_policy: dict[str, Any], updated_policy: dict[str, Any]) -> dict[str, Any]:
    """Assess the risk level of proposed autoscaler policy changes."""
    risks = []

    # Enabling autoscaler from disabled
    if not current_policy.get("enabled") and updated_policy.get("enabled"):
        risks.append({
            "level": "HIGH",
            "change": "Enabling the autoscaler",
            "impact": "CAST.AI will start adding and removing nodes based on cluster demand",
        })

    # Disabling autoscaler
    if current_policy.get("enabled") and not updated_policy.get("enabled"):
        risks.append({
            "level": "MEDIUM",
            "change": "Disabling the autoscaler",
            "impact": "CAST.AI will stop managing cluster nodes - no automatic scaling",
        })

    # Enabling node downscaler
    current_downscaler = current_policy.get("nodeDownscaler", {})
    updated_downscaler = updated_policy.get("nodeDownscaler", {})
    if not current_downscaler.get("enabled") and updated_downscaler.get("enabled"):
        risks.append({
            "level": "HIGH",
            "change": "Enabling node downscaler",
            "impact": "Empty and underutilized nodes will be removed from the cluster",
        })

    # Enabling evictor through this tool
    current_evictor = current_downscaler.get("evictor", {})
    updated_evictor = updated_downscaler.get("evictor", {})
    if not current_evictor.get("enabled") and updated_evictor.get("enabled"):
        if not updated_evictor.get("dryRun"):
            risks.append({
                "level": "HIGH",
                "change": "Enabling evictor (not dry-run)",
                "impact": "Pods will be evicted and rescheduled to optimize node usage",
            })

    # Enabling spot instances
    current_spot = current_policy.get("spotInstances", {})
    updated_spot = updated_policy.get("spotInstances", {})
    if not current_spot.get("enabled") and updated_spot.get("enabled"):
        risks.append({
            "level": "MEDIUM",
            "change": "Enabling spot instances",
            "impact": "New nodes may use spot/preemptible instances which can be reclaimed",
        })

    # Determine overall risk
    risk_levels = [r["level"] for r in risks]
    if "CRITICAL" in risk_levels:
        overall_risk = "CRITICAL"
    elif "HIGH" in risk_levels:
        overall_risk = "HIGH"
    elif "MEDIUM" in risk_levels:
        overall_risk = "MEDIUM"
    else:
        overall_risk = "LOW"

    return {
        "overall_risk": overall_risk,
        "risks": risks,
        "requires_extra_caution": overall_risk in ["CRITICAL", "HIGH"],
    }


def register_autoscaler_tool(mcp: FastMCP):
    """Register the autoscaler policy update tool."""

    @mcp.tool()
    async def update_cluster_autoscaler_policies(
        cluster_id_or_name: str,
        enabled: bool | None = None,
        unschedulable_pods_enabled: bool | None = None,
        headroom_cpu_percentage: int | None = None,
        headroom_memory_percentage: int | None = None,
        headroom_enabled: bool | None = None,
        spot_headroom_cpu_percentage: int | None = None,
        spot_headroom_memory_percentage: int | None = None,
        spot_headroom_enabled: bool | None = None,
        custom_instances_enabled: bool | None = None,
        node_constraints_enabled: bool | None = None,
        node_constraints_min_cpu_cores: int | None = None,
        node_constraints_max_cpu_cores: int | None = None,
        node_constraints_min_ram_mib: int | None = None,
        node_constraints_max_ram_mib: int | None = None,
        cluster_limits_enabled: bool | None = None,
        cluster_limits_cpu_min_cores: int | None = None,
        cluster_limits_cpu_max_cores: int | None = None,
        spot_instances_enabled: bool | None = None,
        spot_diversity_enabled: bool | None = None,
        spot_diversity_price_increase_limit_percent: int | None = None,
        spot_backups_enabled: bool | None = None,
        spot_interruption_predictions_enabled: bool | None = None,
        node_downscaler_enabled: bool | None = None,
        empty_nodes_enabled: bool | None = None,
        empty_nodes_delay_seconds: int | None = None,
        is_scoped_mode: bool | None = None,
        node_templates_partial_matching_enabled: bool | None = None,
        confirm_impact: bool = False,
    ) -> str:
        """
        Update autoscaler policies for a CAST.AI Kubernetes cluster.

        USE THIS TOOL WHEN THE USER WANTS TO:
        - "enable autoscaling" / "turn on autoscaler" / "start autoscaling"
        - "add nodes" / "scale up" / "provision nodes" / "enable upscaling"
        - "remove nodes" / "scale down" / "downscale" / "enable downscaling"
        - "optimize cluster" / "optimize my cluster" / "enable optimization"
        - "reduce costs" / "save money" / "cut cloud spend" / "cost optimization"
        - "take over cluster management" / "let CAST.AI manage my cluster"
        - "enable spot instances" / "use spot nodes" / "use preemptible instances"
        - "set cluster limits" / "limit cluster size" / "cap cluster resources"
        - "add headroom" / "reserve capacity" / "buffer resources"
        - "disable autoscaling" / "stop scaling" / "turn off autoscaler"
        - "enable/disable unschedulable pods policy" / "enable/disable upscaling policy"
        - "stop adding nodes" / "stop removing nodes"
        - Any request about changing how the cluster scales, adds nodes, or removes nodes

        This is the primary tool for controlling CAST.AI cluster autoscaling behavior.
        It manages the policies that determine when and how nodes are added or removed.

        ⚠️  Changes take effect immediately. Set confirm_impact=true to apply.

        WHAT EACH SECTION CONTROLS:
        - **enabled**: Master switch for the entire autoscaler (MUST be true for anything to work)
        - **unschedulable_pods** (= upscaling): Adds new nodes when pods can't be scheduled (MUST be enabled for autoscaler to add nodes)
        - **node_downscaler** (= downscaling): Removes empty/underutilized nodes (MUST be enabled for autoscaler to remove nodes)
        - **cluster_limits**: Caps the total cluster size (CPU cores)
        - **spot_instances**: Uses cheaper spot/preemptible VMs for cost savings

        CRITICAL: The "enabled" flag alone does NOTHING. You MUST also enable sub-policies:
        - To add nodes: enabled=true AND unschedulable_pods_enabled=true
        - To remove nodes: enabled=true AND node_downscaler_enabled=true AND empty_nodes_enabled=true

        REQUIRED CONFIGURATIONS (always use these combinations):
        - "Enable autoscaler" / "enable autoscaling" / "enable optimization": enabled=true, unschedulable_pods_enabled=true, node_downscaler_enabled=true, empty_nodes_enabled=true
        - "Enable upscaling" / "add nodes": enabled=true, unschedulable_pods_enabled=true
        - "Enable downscaling" / "remove nodes": enabled=true, node_downscaler_enabled=true, empty_nodes_enabled=true
        - Cost optimization with spot: enabled=true, unschedulable_pods_enabled=true, node_downscaler_enabled=true, empty_nodes_enabled=true, spot_instances_enabled=true
        - Add headroom buffer: headroom_enabled=true, headroom_cpu_percentage=10, headroom_memory_percentage=10
        - Limit cluster size: cluster_limits_enabled=true, cluster_limits_cpu_max_cores=100
        - Disable everything: enabled=false

        Args:
            cluster_id_or_name: Cluster ID (UUID) or cluster name
            enabled: Master switch - enable/disable entire autoscaler
            unschedulable_pods_enabled: Enable provisioning nodes for unschedulable pods
            headroom_cpu_percentage: Extra on-demand CPU headroom percentage (0-100)
            headroom_memory_percentage: Extra on-demand memory headroom percentage (0-100)
            headroom_enabled: Enable on-demand headroom
            spot_headroom_cpu_percentage: Extra spot CPU headroom percentage (0-100)
            spot_headroom_memory_percentage: Extra spot memory headroom percentage (0-100)
            spot_headroom_enabled: Enable spot headroom
            custom_instances_enabled: Allow custom instance types
            node_constraints_enabled: Enable node size constraints
            node_constraints_min_cpu_cores: Minimum CPU cores per node
            node_constraints_max_cpu_cores: Maximum CPU cores per node
            node_constraints_min_ram_mib: Minimum RAM in MiB per node
            node_constraints_max_ram_mib: Maximum RAM in MiB per node
            cluster_limits_enabled: Enable cluster-wide CPU limits
            cluster_limits_cpu_min_cores: Minimum total CPU cores for cluster
            cluster_limits_cpu_max_cores: Maximum total CPU cores for cluster
            spot_instances_enabled: Enable spot/preemptible instances
            spot_diversity_enabled: Enable spot instance diversity
            spot_diversity_price_increase_limit_percent: Max price increase for diversity (0-100)
            spot_backups_enabled: Enable on-demand backups for spot instances
            spot_interruption_predictions_enabled: Enable spot interruption predictions
            node_downscaler_enabled: Enable node downscaler (removes underutilized nodes)
            empty_nodes_enabled: Enable empty node removal
            empty_nodes_delay_seconds: Delay before removing empty nodes (seconds)
            is_scoped_mode: Enable scoped mode (namespace-specific autoscaling)
            node_templates_partial_matching_enabled: Enable partial matching for node templates
            confirm_impact: REQUIRED - Set to true to confirm you understand the impact

        Returns:
            JSON string with update result including before/after comparison
        """
        try:
            # Step 1: Resolve cluster
            cluster_id = await resolve_cluster_id(cluster_id_or_name)
            logger.info("Starting autoscaler policy update", cluster_id=cluster_id, confirmed=confirm_impact)

            # Step 2: Read current policy
            try:
                current_policy = await make_request("GET", f"/v1/kubernetes/clusters/{cluster_id}/policies")
            except Exception as e:
                return json.dumps({
                    "error": "Failed to read current policy",
                    "details": str(e),
                    "suggestion": "Verify cluster name/ID is correct using list_clusters",
                })

            # Step 3: Build updated policy (read-modify-write)
            updated_policy = json.loads(json.dumps(current_policy))  # deep copy
            changes = {}

            def track_change(key: str, current_val, new_val, target_dict: dict, target_key: str):
                """Only record and apply the change if the value actually differs."""
                target_dict[target_key] = new_val
                if current_val != new_val:
                    changes[key] = {"from": current_val, "to": new_val}

            # Master switch
            if enabled is not None:
                track_change("enabled", current_policy.get("enabled"), enabled, updated_policy, "enabled")

            # Unschedulable pods section
            up = updated_policy.setdefault("unschedulablePods", {})
            cup = current_policy.get("unschedulablePods", {})

            if unschedulable_pods_enabled is not None:
                track_change("unschedulablePods.enabled", cup.get("enabled"), unschedulable_pods_enabled, up, "enabled")

            # On-demand headroom
            headroom = up.setdefault("headroom", {})
            c_headroom = cup.get("headroom", {})

            if headroom_enabled is not None:
                track_change("headroom.enabled", c_headroom.get("enabled"), headroom_enabled, headroom, "enabled")
            if headroom_cpu_percentage is not None:
                track_change("headroom.cpuPercentage", c_headroom.get("cpuPercentage"), headroom_cpu_percentage, headroom, "cpuPercentage")
            if headroom_memory_percentage is not None:
                track_change("headroom.memoryPercentage", c_headroom.get("memoryPercentage"), headroom_memory_percentage, headroom, "memoryPercentage")

            # Spot headroom
            headroom_spot = up.setdefault("headroomSpot", {})
            c_headroom_spot = cup.get("headroomSpot", {})

            if spot_headroom_enabled is not None:
                track_change("headroomSpot.enabled", c_headroom_spot.get("enabled"), spot_headroom_enabled, headroom_spot, "enabled")
            if spot_headroom_cpu_percentage is not None:
                track_change("headroomSpot.cpuPercentage", c_headroom_spot.get("cpuPercentage"), spot_headroom_cpu_percentage, headroom_spot, "cpuPercentage")
            if spot_headroom_memory_percentage is not None:
                track_change("headroomSpot.memoryPercentage", c_headroom_spot.get("memoryPercentage"), spot_headroom_memory_percentage, headroom_spot, "memoryPercentage")

            # Custom instances
            if custom_instances_enabled is not None:
                track_change("customInstancesEnabled", cup.get("customInstancesEnabled"), custom_instances_enabled, up, "customInstancesEnabled")

            # Node constraints
            nc = up.setdefault("nodeConstraints", {})
            c_nc = cup.get("nodeConstraints", {})

            if node_constraints_enabled is not None:
                track_change("nodeConstraints.enabled", c_nc.get("enabled"), node_constraints_enabled, nc, "enabled")
            if node_constraints_min_cpu_cores is not None:
                track_change("nodeConstraints.minCpuCores", c_nc.get("minCpuCores"), node_constraints_min_cpu_cores, nc, "minCpuCores")
            if node_constraints_max_cpu_cores is not None:
                track_change("nodeConstraints.maxCpuCores", c_nc.get("maxCpuCores"), node_constraints_max_cpu_cores, nc, "maxCpuCores")
            if node_constraints_min_ram_mib is not None:
                track_change("nodeConstraints.minRamMib", c_nc.get("minRamMib"), node_constraints_min_ram_mib, nc, "minRamMib")
            if node_constraints_max_ram_mib is not None:
                track_change("nodeConstraints.maxRamMib", c_nc.get("maxRamMib"), node_constraints_max_ram_mib, nc, "maxRamMib")

            # Cluster limits
            cl = updated_policy.setdefault("clusterLimits", {})
            c_cl = current_policy.get("clusterLimits", {})

            if cluster_limits_enabled is not None:
                track_change("clusterLimits.enabled", c_cl.get("enabled"), cluster_limits_enabled, cl, "enabled")

            cl_cpu = cl.setdefault("cpu", {})
            c_cl_cpu = c_cl.get("cpu", {})

            if cluster_limits_cpu_min_cores is not None:
                track_change("clusterLimits.cpu.minCores", c_cl_cpu.get("minCores"), cluster_limits_cpu_min_cores, cl_cpu, "minCores")
            if cluster_limits_cpu_max_cores is not None:
                track_change("clusterLimits.cpu.maxCores", c_cl_cpu.get("maxCores"), cluster_limits_cpu_max_cores, cl_cpu, "maxCores")

            # Spot instances
            si = updated_policy.setdefault("spotInstances", {})
            c_si = current_policy.get("spotInstances", {})

            if spot_instances_enabled is not None:
                track_change("spotInstances.enabled", c_si.get("enabled"), spot_instances_enabled, si, "enabled")
            if spot_diversity_enabled is not None:
                track_change("spotInstances.spotDiversityEnabled", c_si.get("spotDiversityEnabled"), spot_diversity_enabled, si, "spotDiversityEnabled")
            if spot_diversity_price_increase_limit_percent is not None:
                track_change(
                    "spotInstances.spotDiversityPriceIncreaseLimitPercent",
                    c_si.get("spotDiversityPriceIncreaseLimitPercent"),
                    spot_diversity_price_increase_limit_percent,
                    si, "spotDiversityPriceIncreaseLimitPercent",
                )

            sb = si.setdefault("spotBackups", {})
            c_sb = c_si.get("spotBackups", {})

            if spot_backups_enabled is not None:
                track_change("spotInstances.spotBackups.enabled", c_sb.get("enabled"), spot_backups_enabled, sb, "enabled")

            sip = si.setdefault("spotInterruptionPredictions", {})
            c_sip = c_si.get("spotInterruptionPredictions", {})

            if spot_interruption_predictions_enabled is not None:
                track_change(
                    "spotInstances.spotInterruptionPredictions.enabled",
                    c_sip.get("enabled"),
                    spot_interruption_predictions_enabled,
                    sip, "enabled",
                )

            # Node downscaler
            nd = updated_policy.setdefault("nodeDownscaler", {})
            c_nd = current_policy.get("nodeDownscaler", {})

            if node_downscaler_enabled is not None:
                track_change("nodeDownscaler.enabled", c_nd.get("enabled"), node_downscaler_enabled, nd, "enabled")

            en = nd.setdefault("emptyNodes", {})
            c_en = c_nd.get("emptyNodes", {})

            if empty_nodes_enabled is not None:
                track_change("emptyNodes.enabled", c_en.get("enabled"), empty_nodes_enabled, en, "enabled")
            if empty_nodes_delay_seconds is not None:
                track_change("emptyNodes.delaySeconds", c_en.get("delaySeconds"), empty_nodes_delay_seconds, en, "delaySeconds")

            # Top-level flags
            if is_scoped_mode is not None:
                track_change("isScopedMode", current_policy.get("isScopedMode"), is_scoped_mode, updated_policy, "isScopedMode")

            if node_templates_partial_matching_enabled is not None:
                track_change(
                    "nodeTemplatesPartialMatchingEnabled",
                    current_policy.get("nodeTemplatesPartialMatchingEnabled"),
                    node_templates_partial_matching_enabled,
                    updated_policy, "nodeTemplatesPartialMatchingEnabled",
                )

            # Step 4: Check if any parameters were provided at all
            any_param_set = any(v is not None for v in [
                enabled, unschedulable_pods_enabled, headroom_cpu_percentage,
                headroom_memory_percentage, headroom_enabled, spot_headroom_cpu_percentage,
                spot_headroom_memory_percentage, spot_headroom_enabled, custom_instances_enabled,
                node_constraints_enabled, node_constraints_min_cpu_cores, node_constraints_max_cpu_cores,
                node_constraints_min_ram_mib, node_constraints_max_ram_mib, cluster_limits_enabled,
                cluster_limits_cpu_min_cores, cluster_limits_cpu_max_cores, spot_instances_enabled,
                spot_diversity_enabled, spot_diversity_price_increase_limit_percent,
                spot_backups_enabled, spot_interruption_predictions_enabled, node_downscaler_enabled,
                empty_nodes_enabled, empty_nodes_delay_seconds, is_scoped_mode,
                node_templates_partial_matching_enabled,
            ])
            if not any_param_set:
                return json.dumps({
                    "error": "No changes specified",
                    "suggestion": "Provide at least one parameter to update",
                    "current_config_summary": {
                        "enabled": current_policy.get("enabled"),
                        "unschedulablePods": cup.get("enabled"),
                        "nodeDownscaler": c_nd.get("enabled"),
                        "emptyNodes": c_en.get("enabled"),
                        "clusterLimits": c_cl.get("enabled"),
                        "spotInstances": c_si.get("enabled"),
                    },
                })

            # Step 5: Assess risk
            risk_assessment = assess_autoscaler_risk(current_policy, updated_policy)

            # Step 6: Require confirmation
            if not confirm_impact:
                return json.dumps({
                    "error": "Confirmation required",
                    "message": (
                        "⚠️  SAFETY CHECK: You must set confirm_impact=true to proceed.\n\n"
                        "This operation will modify cluster autoscaling behavior. "
                        "Changes take effect immediately and affect how nodes are added/removed."
                    ),
                    "proposed_changes": changes,
                    "risk_assessment": risk_assessment,
                    "required_action": "Set confirm_impact=true to acknowledge and proceed",
                })

            if risk_assessment["requires_extra_caution"]:
                logger.warning(
                    "High-risk autoscaler policy update confirmed",
                    cluster_id=cluster_id,
                    risk_level=risk_assessment["overall_risk"],
                )

            # Step 7: Apply update
            logger.info("Updating autoscaler policies", cluster_id=cluster_id, changes=len(changes))
            try:
                result = await make_request("PUT", f"/v1/kubernetes/clusters/{cluster_id}/policies", body=updated_policy)
            except Exception as e:
                logger.error("Autoscaler policy update failed", error=str(e))
                return json.dumps({
                    "error": "Failed to update policy",
                    "details": str(e),
                    "rollback_suggestion": "No changes were applied. Your previous configuration is still active.",
                })

            # Step 8: Verify
            logger.info("Verifying autoscaler policy update", cluster_id=cluster_id)
            try:
                verified_policy = await make_request("GET", f"/v1/kubernetes/clusters/{cluster_id}/policies")
            except Exception:
                verified_policy = None

            # Step 9: Return success
            logger.info("Autoscaler policy update completed", cluster_id=cluster_id, changes=len(changes))

            response = {
                "status": "success",
                "message": "Autoscaler policies updated successfully",
                "cluster_id": cluster_id,
                "changes_applied": changes if changes else "All values were already at desired state (applied anyway to ensure consistency)",
                "risk_assessment": risk_assessment,
                "verified_config_summary": {
                    "enabled": (verified_policy or updated_policy).get("enabled"),
                    "unschedulablePods": (verified_policy or updated_policy).get("unschedulablePods", {}).get("enabled"),
                    "nodeDownscaler": (verified_policy or updated_policy).get("nodeDownscaler", {}).get("enabled"),
                    "clusterLimits": (verified_policy or updated_policy).get("clusterLimits", {}).get("enabled"),
                    "spotInstances": (verified_policy or updated_policy).get("spotInstances", {}).get("enabled"),
                },
                "timestamp": result.get("updatedAt") if isinstance(result, dict) else None,
            }

            if risk_assessment["overall_risk"] in ["CRITICAL", "HIGH"]:
                response["monitoring_recommendation"] = (
                    "⚠️  Monitor your cluster closely for the next 15-30 minutes. "
                    "Watch for unexpected node additions/removals or pod disruptions. "
                    "If issues occur, set enabled=false to disable the autoscaler."
                )

            response["rollback"] = {
                "disable_autoscaler": "Set enabled=false",
                "disable_downscaler": "Set node_downscaler_enabled=false",
                "disable_spot": "Set spot_instances_enabled=false",
            }

            return json.dumps(response, indent=2)

        except Exception as e:
            logger.error("Autoscaler policy update error", error=str(e), cluster_id_or_name=cluster_id_or_name)
            return json.dumps({
                "error": str(e),
                "troubleshooting": {
                    "check_cluster": "Verify cluster name/ID using list_clusters",
                    "check_permissions": "Ensure API key has write permissions",
                    "check_api_status": "Verify CAST.AI API is accessible",
                },
            })
