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

Provides tools to generate, inspect, and execute cluster rebalancing plans.
Rebalancing optimizes node placement by replacing existing nodes with more
cost-effective ones while respecting workload constraints.
"""

import asyncio
import json

from fastmcp import FastMCP

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


def _summarize_plan(plan: dict) -> dict:
    """Extract a human-readable summary from a rebalancing plan response."""
    configs = plan.get("configurations", {})
    blue = configs.get("blue", {})
    green = configs.get("green", {})
    diff = configs.get("diff", {})
    achieved_diff = configs.get("achievedDiff", {})

    blue_totals = blue.get("totals", {})
    green_totals = green.get("totals", {})

    summary = {
        "rebalancing_plan_id": plan.get("rebalancingPlanId"),
        "status": plan.get("status"),
        "created_at": plan.get("createdAt"),
        "updated_at": plan.get("updatedAt"),
        "min_nodes": plan.get("minNodes"),
        "evict_gracefully": plan.get("evictGracefully"),
    }

    # Current state (blue)
    summary["current_nodes"] = {
        "count": len(blue.get("nodes", [])),
        "total_cpu_milli": blue_totals.get("milliCpu"),
        "total_memory_mib": blue_totals.get("memoryMib"),
        "monthly_cost": blue_totals.get("priceMonthly"),
        "hourly_cost": blue_totals.get("priceHourly"),
    }

    # Proposed state (green)
    summary["proposed_nodes"] = {
        "count": len(green.get("nodes", [])),
        "total_cpu_milli": green_totals.get("milliCpu"),
        "total_memory_mib": green_totals.get("memoryMib"),
        "monthly_cost": green_totals.get("priceMonthly"),
        "hourly_cost": green_totals.get("priceHourly"),
    }

    # Savings
    summary["estimated_savings"] = {
        "hourly": diff.get("priceHourly"),
        "monthly": diff.get("priceMonthly"),
        "savings_percentage": diff.get("savingsPercentage"),
        "cluster_savings_percentage": diff.get("clusterSavingsPercentage"),
    }

    if achieved_diff:
        summary["achieved_savings"] = {
            "hourly": achieved_diff.get("priceHourly"),
            "monthly": achieved_diff.get("priceMonthly"),
            "savings_percentage": achieved_diff.get("savingsPercentage"),
        }

    # Errors
    errors = plan.get("errors", [])
    if errors:
        summary["errors"] = errors

    return summary


def register_rebalancing_tools(mcp: FastMCP):
    """Register rebalancing-related MCP tools."""

    @mcp.tool()
    async def generate_rebalancing_plan(
        cluster_id_or_name: str,
        min_nodes: int = 3,
        evict_gracefully: bool = True,
        node_ids: str | None = None,
    ) -> str:
        """
        Generate a rebalancing plan for a CAST.AI Kubernetes cluster.

        USE THIS TOOL WHEN THE USER WANTS TO:
        - "rebalance my cluster" / "rebalance nodes"
        - "optimize node placement" / "optimize nodes"
        - "generate rebalancing plan" / "create rebalancing plan"
        - "reduce node costs" / "find cheaper nodes"
        - "consolidate nodes" / "right-size nodes"
        - "replace nodes with better ones" / "get better instance types"
        - "how much can I save by rebalancing?"

        This generates a plan that shows what nodes would be replaced and the
        estimated cost savings. The plan is NOT executed automatically — you must
        call execute_rebalancing_plan to apply it.

        WORKFLOW:
        1. generate_rebalancing_plan → creates plan, returns savings estimate
        2. Review the plan with the user (savings, node changes)
        3. execute_rebalancing_plan → applies the plan (requires confirmation)

        Args:
            cluster_id_or_name: Cluster ID (UUID) or cluster name
            min_nodes: Minimum number of nodes to keep after rebalancing (default: 3)
            evict_gracefully: If true, nodes that fail to drain are kept instead of force-deleted (default: true)
            node_ids: Optional comma-separated list of specific node IDs to rebalance. If empty, all nodes are considered.

        Returns:
            JSON string with the rebalancing plan summary including estimated savings
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)
            logger.info("Generating rebalancing plan", cluster_id=cluster_id, min_nodes=min_nodes)

            # Build request body
            body: dict = {
                "minNodes": min_nodes,
                "evictGracefully": evict_gracefully,
            }

            if node_ids:
                body["rebalancingNodes"] = [{"nodeId": nid.strip()} for nid in node_ids.split(",")]

            # Generate plan
            try:
                result = await make_request(
                    "POST",
                    f"/v1/kubernetes/clusters/{cluster_id}/rebalancing-plans",
                    body=body,
                )
            except Exception as e:
                return json.dumps({
                    "error": "Failed to generate rebalancing plan",
                    "details": str(e),
                    "suggestion": "Verify the cluster has active nodes and autoscaler policies are enabled",
                })

            plan_id = result.get("rebalancingPlanId")
            if not plan_id:
                return json.dumps({
                    "error": "No plan ID returned",
                    "details": "The API did not return a rebalancing plan ID",
                    "raw_response": result,
                })

            logger.info("Rebalancing plan created, waiting for generation", plan_id=plan_id)

            # Poll until plan leaves "generating" status (max ~60 seconds)
            plan = None
            for _ in range(30):
                await asyncio.sleep(2)
                try:
                    plan = await make_request(
                        "GET",
                        f"/v1/kubernetes/clusters/{cluster_id}/rebalancing-plans/{plan_id}",
                        params={"includeConfigurations": "true"},
                    )
                except Exception:
                    continue

                status = plan.get("status", "")
                if status != "generating":
                    break

            if not plan:
                return json.dumps({
                    "error": "Could not retrieve plan after generation",
                    "rebalancing_plan_id": plan_id,
                    "suggestion": "Use get_rebalancing_plan to check the plan status manually",
                })

            summary = _summarize_plan(plan)

            if plan.get("status") == "generating":
                summary["message"] = (
                    "Plan is still being generated. Use get_rebalancing_plan to check progress."
                )
            elif plan.get("status") == "generated":
                summary["message"] = (
                    "Plan is ready! Review the estimated savings above. "
                    "To apply this plan, use execute_rebalancing_plan with the plan ID."
                )
                summary["next_step"] = f"execute_rebalancing_plan(cluster_id_or_name='{cluster_id_or_name}', rebalancing_plan_id='{plan_id}', confirm_impact=True)"
            elif plan.get("status") == "error":
                summary["message"] = "Plan generation failed. Check the errors field for details."
            else:
                summary["message"] = f"Plan status: {plan.get('status')}"

            logger.info("Rebalancing plan generated", plan_id=plan_id, status=plan.get("status"))
            return json.dumps(summary, indent=2)

        except Exception as e:
            logger.error("Rebalancing plan generation error", error=str(e))
            return json.dumps({
                "error": str(e),
                "troubleshooting": {
                    "check_cluster": "Verify cluster name/ID using list_clusters",
                    "check_policies": "Ensure autoscaler policies are enabled",
                    "check_permissions": "Ensure API key has write permissions",
                },
            })

    @mcp.tool()
    async def get_rebalancing_plan(
        cluster_id_or_name: str,
        rebalancing_plan_id: str,
    ) -> str:
        """
        Get the status and details of a rebalancing plan.

        USE THIS TOOL WHEN THE USER WANTS TO:
        - "check rebalancing status" / "how is rebalancing going?"
        - "get rebalancing plan details" / "show me the plan"
        - "is rebalancing done?" / "rebalancing progress"
        - "what happened with the rebalancing?"

        Returns the current status, node changes, savings estimate, and any errors.

        Plan statuses: generating → generated → creating_nodes → preparing_nodes →
        draining_nodes → deleting_nodes → finished/partially_finished/error

        Args:
            cluster_id_or_name: Cluster ID (UUID) or cluster name
            rebalancing_plan_id: The rebalancing plan ID (from generate_rebalancing_plan)

        Returns:
            JSON string with plan status, savings, and node details
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)
            logger.info("Getting rebalancing plan", cluster_id=cluster_id, plan_id=rebalancing_plan_id)

            plan = await make_request(
                "GET",
                f"/v1/kubernetes/clusters/{cluster_id}/rebalancing-plans/{rebalancing_plan_id}",
                params={"includeConfigurations": "true", "includeOperations": "true"},
            )

            summary = _summarize_plan(plan)

            # Add operation progress if executing
            operations = plan.get("operations", [])
            if operations:
                op_summary = {"total": len(operations), "by_type": {}, "by_status": {"completed": 0, "pending": 0, "failed": 0}}
                for op in operations:
                    op_type = op.get("type", "unknown")
                    op_summary["by_type"][op_type] = op_summary["by_type"].get(op_type, 0) + 1
                    if op.get("finishedAt"):
                        if op.get("error"):
                            op_summary["by_status"]["failed"] += 1
                        else:
                            op_summary["by_status"]["completed"] += 1
                    else:
                        op_summary["by_status"]["pending"] += 1
                summary["operations"] = op_summary

            # Status-specific messages
            status = plan.get("status", "")
            if status == "generated":
                summary["message"] = "Plan is ready to execute. Use execute_rebalancing_plan to apply it."
            elif status == "finished":
                summary["message"] = "Rebalancing completed successfully!"
            elif status == "partially_finished":
                summary["message"] = "Rebalancing partially completed. Some operations may have failed."
            elif status == "error":
                summary["message"] = "Rebalancing failed. Check the errors field for details."
            elif status in ("creating_nodes", "preparing_nodes", "draining_nodes", "deleting_nodes"):
                summary["message"] = f"Rebalancing is in progress (phase: {status})."

            return json.dumps(summary, indent=2)

        except Exception as e:
            logger.error("Get rebalancing plan error", error=str(e))
            return json.dumps({"error": str(e)})

    @mcp.tool()
    async def execute_rebalancing_plan(
        cluster_id_or_name: str,
        rebalancing_plan_id: str,
        confirm_impact: bool = False,
    ) -> str:
        """
        Execute a previously generated rebalancing plan.

        USE THIS TOOL WHEN THE USER WANTS TO:
        - "execute rebalancing" / "apply rebalancing plan" / "start rebalancing"
        - "go ahead with rebalancing" / "proceed with the plan"
        - "run the rebalancing" / "do the rebalancing"
        - User confirms they want to apply a previously generated plan

        ⚠️  WARNING: This will replace nodes in the cluster. Workloads will be
        drained and rescheduled onto new, more cost-effective nodes.

        PREREQUISITES:
        - A plan must be generated first using generate_rebalancing_plan
        - The plan must be in "generated" status
        - Set confirm_impact=true to proceed

        Args:
            cluster_id_or_name: Cluster ID (UUID) or cluster name
            rebalancing_plan_id: The rebalancing plan ID (from generate_rebalancing_plan)
            confirm_impact: REQUIRED - Set to true to confirm you understand nodes will be replaced

        Returns:
            JSON string with execution status and monitoring instructions
        """
        try:
            cluster_id = await resolve_cluster_id(cluster_id_or_name)
            logger.info("Execute rebalancing plan requested", cluster_id=cluster_id, plan_id=rebalancing_plan_id, confirmed=confirm_impact)

            # Fetch plan to show current state before execution
            try:
                plan = await make_request(
                    "GET",
                    f"/v1/kubernetes/clusters/{cluster_id}/rebalancing-plans/{rebalancing_plan_id}",
                    params={"includeConfigurations": "true"},
                )
            except Exception as e:
                return json.dumps({
                    "error": "Failed to fetch rebalancing plan",
                    "details": str(e),
                    "suggestion": "Verify the plan ID is correct",
                })

            status = plan.get("status", "")

            # Validate plan is executable
            if status != "generated":
                return json.dumps({
                    "error": f"Plan cannot be executed - current status is '{status}'",
                    "suggestion": "Only plans with status 'generated' can be executed",
                    "current_status": status,
                })

            summary = _summarize_plan(plan)

            # 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 will replace nodes in your cluster. Workloads will be drained "
                        "and rescheduled onto new nodes. The process takes several minutes."
                    ),
                    "plan_summary": summary,
                    "required_action": "Set confirm_impact=true to acknowledge and proceed",
                })

            # Execute
            logger.info("Executing rebalancing plan", cluster_id=cluster_id, plan_id=rebalancing_plan_id)
            try:
                result = await make_request(
                    "POST",
                    f"/v1/kubernetes/clusters/{cluster_id}/rebalancing-plans/{rebalancing_plan_id}/execute",
                )
            except Exception as e:
                logger.error("Rebalancing execution failed", error=str(e))
                return json.dumps({
                    "error": "Failed to execute rebalancing plan",
                    "details": str(e),
                    "rollback_suggestion": "The plan was not executed. Your cluster is unchanged.",
                })

            # Return execution confirmation
            exec_summary = _summarize_plan(result) if isinstance(result, dict) and result.get("status") else summary
            exec_summary["execution_status"] = "started"
            exec_summary["message"] = (
                "Rebalancing execution started! The process will:\n"
                "1. Create new optimized nodes\n"
                "2. Prepare new nodes (install components)\n"
                "3. Drain workloads from old nodes\n"
                "4. Delete old nodes\n\n"
                "Use get_rebalancing_plan to monitor progress."
            )
            exec_summary["monitoring"] = {
                "check_progress": f"get_rebalancing_plan(cluster_id_or_name='{cluster_id_or_name}', rebalancing_plan_id='{rebalancing_plan_id}')",
                "recommendation": "Check progress every 2-5 minutes. Full rebalancing typically takes 10-30 minutes.",
            }

            logger.info("Rebalancing execution started", cluster_id=cluster_id, plan_id=rebalancing_plan_id)
            return json.dumps(exec_summary, indent=2)

        except Exception as e:
            logger.error("Rebalancing execution error", error=str(e))
            return json.dumps({
                "error": str(e),
                "troubleshooting": {
                    "check_cluster": "Verify cluster name/ID using list_clusters",
                    "check_plan": "Verify plan ID using get_rebalancing_plan",
                    "check_permissions": "Ensure API key has write permissions",
                },
            })
