"""
Graph Partitioner - Memory-aware partitioning for offloaded execution.

Given a ModelGraph and a memory budget, produces an execution plan that
specifies:
- Which layers/components to keep resident vs. offload
- Loading order and prefetch schedule
- Memory allocation per step
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Optional

from dokodemo_ai.graph.compiler import LayerInfo, ModelGraph

logger = logging.getLogger(__name__)


@dataclass
class LoadAction:
    """A single load action in the execution plan."""

    layer_index: int
    component_name: str  # Name of the component to load
    component_kind: str  # "attention", "mlp", "router", "expert", "norm", etc.
    size_bytes: int
    expert_index: Optional[int] = None  # For expert components
    priority: float = 1.0  # Higher = more important to cache


@dataclass
class StepPlan:
    """Execution plan for a single layer during one forward pass."""

    layer_index: int
    is_moe: bool
    preload_actions: list[LoadAction] = field(default_factory=list)
    compute_actions: list[str] = field(default_factory=list)
    evict_after: bool = True  # Whether to evict after computation
    keep_resident: bool = False  # Whether to keep in cache


@dataclass
class ExecutionPlan:
    """Complete execution plan for one forward pass through all layers."""

    steps: list[StepPlan]
    memory_budget_bytes: int
    peak_memory_bytes: int
    resident_layers: list[int]  # Layers to keep in GPU/CPU memory

    def summary(self) -> str:
        total_load = sum(
            sum(a.size_bytes for a in step.preload_actions)
            for step in self.steps
        )
        resident_load = sum(
            sum(a.size_bytes for a in step.preload_actions)
            for step in self.steps
            if step.keep_resident
        )
        lines = [
            f"Execution plan: {len(self.steps)} steps",
            f"Memory budget: {self.memory_budget_bytes / 1e9:.2f} GB",
            f"Peak memory: {self.peak_memory_bytes / 1e9:.2f} GB",
            f"Total I/O per token: {total_load / 1e9:.2f} GB",
            f"Resident: {resident_load / 1e9:.2f} GB "
            f"({len(self.resident_layers)} layers)",
        ]
        return "\n".join(lines)


class GraphPartitioner:
    """
    Creates an execution plan from a ModelGraph and memory constraints.

    The partitioner decides:
    1. Which layers can stay resident in memory (hot cache)
    2. The loading schedule for offloaded layers
    3. Prefetch order to maximize I/O overlap
    """

    def __init__(self, memory_budget_bytes: int):
        """
        Args:
            memory_budget_bytes: Available GPU memory for model weights.
                Excludes memory needed for activations, KV cache, etc.
        """
        self.memory_budget_bytes = memory_budget_bytes

    def partition(self, graph: ModelGraph) -> ExecutionPlan:
        """Create an execution plan for the given model graph."""

        # Step 1: Determine which layers can be kept resident
        resident_layers = self._select_resident_layers(graph)

        # Step 2: Create step plans for each layer
        steps = []
        peak_memory = 0
        current_memory = 0

        # Account for embedding and head (always loaded)
        base_memory = graph.embedding.size_bytes + graph.head.size_bytes
        current_memory = base_memory

        for layer in graph.layers:
            step = self._plan_layer(layer, layer.index in resident_layers)
            steps.append(step)

            step_memory = current_memory + sum(
                a.size_bytes for a in step.preload_actions
            )
            peak_memory = max(peak_memory, step_memory)

        plan = ExecutionPlan(
            steps=steps,
            memory_budget_bytes=self.memory_budget_bytes,
            peak_memory_bytes=peak_memory,
            resident_layers=resident_layers,
        )

        logger.info("Created execution plan:\n%s", plan.summary())
        return plan

    def _select_resident_layers(self, graph: ModelGraph) -> list[int]:
        """
        Select layers to keep permanently in GPU memory.

        Strategy: Fill available memory starting with the smallest layers
        (norms, embedding) and first/last layers (which are typically
        most important for output quality).
        """
        available = self.memory_budget_bytes

        # Reserve space for embedding + head (always resident)
        available -= graph.embedding.size_bytes
        available -= graph.head.size_bytes

        if available <= 0:
            # Not enough memory even for embedding + head at full precision.
            # The engine will need to offload everything.
            logger.warning(
                "Memory budget (%.1f GB) is tight. "
                "All layers will be offloaded.",
                self.memory_budget_bytes / 1e9,
            )
            return []

        # Score layers by importance (first and last layers score highest)
        n = len(graph.layers)
        scored: list[tuple[float, int, int]] = []
        for layer in graph.layers:
            # U-shaped importance: first and last layers are most important
            distance_from_edge = min(layer.index, n - 1 - layer.index)
            importance = 1.0 / (1.0 + distance_from_edge)

            # For MoE layers, only consider the non-expert portion for
            # residency (experts are loaded on demand)
            if layer.is_moe:
                size = layer.non_expert_size_bytes
            else:
                size = layer.size_bytes

            scored.append((importance, size, layer.index))

        # Sort by importance descending
        scored.sort(key=lambda x: x[0], reverse=True)

        resident = []
        for importance, size, idx in scored:
            if size <= available:
                resident.append(idx)
                available -= size

        resident.sort()
        return resident

    def _plan_layer(
        self, layer: LayerInfo, keep_resident: bool
    ) -> StepPlan:
        """Create a step plan for a single layer."""
        preload_actions: list[LoadAction] = []

        if layer.is_moe:
            # For MoE layers:
            # - Always load attention + norms + router
            # - Experts are loaded dynamically based on router output

            for comp in layer.components:
                if comp.kind == "expert":
                    # Experts are loaded on-demand by the scheduler
                    continue

                preload_actions.append(LoadAction(
                    layer_index=layer.index,
                    component_name=comp.name,
                    component_kind=comp.kind,
                    size_bytes=comp.size_bytes,
                    priority=2.0 if comp.kind == "router" else 1.0,
                ))

            compute_actions = [
                "load_norms",
                "compute_attention",
                "run_router",
                "load_selected_experts",
                "compute_experts",
                "combine_expert_outputs",
            ]

        else:
            # For dense layers: load everything
            for comp in layer.components:
                preload_actions.append(LoadAction(
                    layer_index=layer.index,
                    component_name=comp.name,
                    component_kind=comp.kind,
                    size_bytes=comp.size_bytes,
                ))

            compute_actions = [
                "load_norms",
                "compute_attention",
                "compute_mlp",
            ]

        return StepPlan(
            layer_index=layer.index,
            is_moe=layer.is_moe,
            preload_actions=preload_actions,
            compute_actions=compute_actions,
            evict_after=not keep_resident,
            keep_resident=keep_resident,
        )
