"""Controller-side logic for pushing updates to agents."""
from __future__ import annotations

import asyncio
import hashlib
import io
import logging
import shutil
import tarfile
import tempfile
from pathlib import Path
from typing import Dict, Optional

import aiohttp

from mesh.updates.manager import UpdateManager
from mesh.updates.models import UpdateStatus

logger = logging.getLogger(__name__)

# Files/dirs to exclude from agent packages (controller-only)
_AGENT_EXCLUDE = {
    "mesh/controller",
    "mesh/dashboard",
    "mesh/updates/server.py",
    "mesh/updates/distributor.py",
    "optimization",
    "data/jepa_policy_model.pt",
    "data/backups",
}


class UpdateDistributor:
    """Pushes update packages from the controller to agents."""

    def __init__(
        self,
        coordinator,
        update_manager: UpdateManager,
        rolling_delay_s: float = 30.0,
    ):
        self.coordinator = coordinator
        self.manager = update_manager
        self.rolling_delay_s = rolling_delay_s

    async def push_to_agent(
        self, node_id: str, package_path: Path
    ) -> UpdateStatus:
        """Upload an update package to a single agent and trigger apply.

        Returns the update status reported by the agent.
        """
        node = self.coordinator.get_node(node_id)
        if node is None:
            logger.error("Node %s not found", node_id)
            return UpdateStatus.FAILED

        agent_url = node.agent_url
        sha256 = UpdateManager.compute_sha256(package_path)

        try:
            url = f"{agent_url}/update/apply"
            timeout = aiohttp.ClientTimeout(total=120)

            with open(package_path, "rb") as f:
                data = aiohttp.FormData()
                data.add_field(
                    "package",
                    f,
                    filename=package_path.name,
                    content_type="application/gzip",
                )
                data.add_field("sha256", sha256)

                async with aiohttp.ClientSession(timeout=timeout) as session:
                    async with session.post(url, data=data) as resp:
                        result = await resp.json()

            if resp.status == 200 and result.get("status") == "ok":
                logger.info("Update pushed to %s (%s)", node.hostname, node_id)
                return UpdateStatus.COMPLETE
            else:
                error = result.get("message", str(result))
                logger.error("Update push to %s failed: %s", node.hostname, error)
                return UpdateStatus.FAILED

        except Exception as e:
            logger.error("Failed to push update to %s: %s", node.hostname, e)
            return UpdateStatus.FAILED

    async def push_to_all(
        self, package_path: Path, rolling: bool = True
    ) -> Dict[str, UpdateStatus]:
        """Push an update to all online agents.

        If rolling=True, updates one at a time with health checks between each.
        If rolling=False, updates all in parallel.

        Returns a mapping of node_id -> status.
        """
        nodes = self.coordinator.get_online_nodes()
        results: Dict[str, UpdateStatus] = {}

        if not nodes:
            logger.warning("No online nodes to update")
            return results

        if rolling:
            results = await self._rolling_update(nodes, package_path)
        else:
            tasks = {
                node.node_id: self.push_to_agent(node.node_id, package_path)
                for node in nodes
            }
            for node_id, coro in tasks.items():
                results[node_id] = await coro

        return results

    async def _rolling_update(
        self, nodes, package_path: Path
    ) -> Dict[str, UpdateStatus]:
        """Update agents one at a time, checking health after each."""
        results: Dict[str, UpdateStatus] = {}

        for i, node in enumerate(nodes):
            logger.info(
                "Rolling update %d/%d: %s (%s)",
                i + 1, len(nodes), node.hostname, node.node_id,
            )

            status = await self.push_to_agent(node.node_id, package_path)
            results[node.node_id] = status

            if status != UpdateStatus.COMPLETE:
                logger.error(
                    "Rolling update stopped: %s failed. "
                    "Remaining nodes not updated.",
                    node.hostname,
                )
                # Mark remaining as not attempted
                for remaining in nodes[i + 1:]:
                    results[remaining.node_id] = UpdateStatus.FAILED
                break

            # Wait and verify health
            if i < len(nodes) - 1:
                logger.info("Waiting %.0fs before next node...", self.rolling_delay_s)
                await asyncio.sleep(self.rolling_delay_s)

                healthy = await self._check_agent_health(node.agent_url)
                if not healthy:
                    logger.error(
                        "Node %s unhealthy after update — stopping rolling update",
                        node.hostname,
                    )
                    for remaining in nodes[i + 1:]:
                        results[remaining.node_id] = UpdateStatus.FAILED
                    break

        return results

    async def _check_agent_health(self, agent_url: str) -> bool:
        """Check if an agent is healthy after update."""
        url = f"{agent_url}/health"
        try:
            timeout = aiohttp.ClientTimeout(total=10)
            async with aiohttp.ClientSession(timeout=timeout) as session:
                async with session.get(url) as resp:
                    return resp.status == 200
        except Exception:
            return False

    def create_agent_package(self, full_package_path: Path) -> Path:
        """Create a smaller agent-only package from a full update package.

        Strips controller, dashboard, JEPA, and other controller-only files.
        Returns the path to the agent package.
        """
        tmp_dir = Path(tempfile.mkdtemp(prefix="mesh-agent-pkg-"))
        agent_pkg = tmp_dir / full_package_path.name.replace(".tar.gz", "-agent.tar.gz")

        with tarfile.open(full_package_path, "r:*") as src:
            with tarfile.open(agent_pkg, "w:gz") as dst:
                for member in src.getmembers():
                    # Skip controller-only files
                    skip = False
                    for exclude in _AGENT_EXCLUDE:
                        if member.name.startswith(exclude) or member.name.lstrip("./").startswith(exclude):
                            skip = True
                            break
                    if skip:
                        continue

                    # Copy member to agent package
                    if member.isfile():
                        f = src.extractfile(member)
                        if f:
                            dst.addfile(member, f)
                    else:
                        dst.addfile(member)

        logger.info(
            "Agent package created: %s (%.1f MB)",
            agent_pkg.name,
            agent_pkg.stat().st_size / 1048576,
        )
        return agent_pkg
