from __future__ import annotations

import asyncio
import importlib.util
import logging
import os
import tempfile
import time
from collections import OrderedDict
from dataclasses import dataclass

import torch

from .aggregator import average_weights
from .job import Job
from .trace import Tracer, TracingConnection
from .transport import Connection, Transport

log = logging.getLogger(__name__)


@dataclass
class Node:
    rank: int
    vast_id: int | None = None
    host: str | None = None
    port: int | None = None
    conn: Connection | None = None
    status: str = "pending"  # pending | ready | running | dead | done


class Scheduler:
    def __init__(self, job: Job, transport: Transport, tracer: Tracer | None = None):
        self.job = job
        self.transport = transport
        self.tracer = tracer
        self.nodes: list[Node] = []
        self.work_dir = tempfile.mkdtemp(prefix="sched_")

    # ------------------------------------------------------------------
    # helpers
    # ------------------------------------------------------------------

    def _active_nodes(self) -> list[Node]:
        return [n for n in self.nodes if n.conn and n.status not in ("dead", "done")]

    def _load_job_module(self):
        spec = importlib.util.spec_from_file_location(
            "job_module", os.path.abspath(self.job.script)
        )
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        return mod

    def _wrap_connections(self):
        """Wrap node connections with TracingConnection if tracer is set."""
        if not self.tracer:
            return
        for node in self.nodes:
            if node.conn and not isinstance(node.conn, TracingConnection):
                node.conn = TracingConnection(node.conn, node.rank, self.tracer)

    # ------------------------------------------------------------------
    # lifecycle phases
    # ------------------------------------------------------------------

    async def provision(self):
        """Rent instances from vast.ai and wait for SSH. Populates self.nodes."""
        from . import vastai

        offers = await vastai.search_offers(self.job.min_vram, limit=self.job.num_nodes)
        if len(offers) < self.job.num_nodes:
            raise RuntimeError(
                f"Need {self.job.num_nodes} instances, only {len(offers)} offers"
            )

        for rank, offer in enumerate(offers[: self.job.num_nodes]):
            await vastai.rent_instance(offer["id"], self.job.docker_image)
            self.nodes.append(Node(rank=rank, vast_id=offer["id"], status="provisioning"))

        # poll until SSH ready (up to 5 min)
        for _ in range(60):
            instances = await vastai.get_instances()
            by_id = {i["id"]: i for i in instances}
            all_ready = True
            for node in self.nodes:
                inst = by_id.get(node.vast_id)
                if inst and inst.get("ssh_host"):
                    node.host = inst["ssh_host"]
                    node.port = inst["ssh_port"]
                    node.status = "ready"
                else:
                    all_ready = False
            if all_ready:
                break
            await asyncio.sleep(5)

        # open connections
        for node in self.nodes:
            if node.host:
                node.conn = await self.transport.connect(node.host, node.port)

        self._wrap_connections()

        ready = self._active_nodes()
        if not ready:
            raise RuntimeError("No nodes became ready")
        log.info(f"Provisioned {len(ready)}/{self.job.num_nodes} nodes")

    async def deploy(self):
        """Push worker script and job module to all active nodes.

        Skips pushing worker.py if it already exists on the remote
        (e.g. baked into a Docker image). Always pushes the job module.
        """
        self._wrap_connections()

        worker_src = os.path.join(
            os.path.dirname(__file__), os.pardir, "worker", "worker.py"
        )
        worker_src = os.path.abspath(worker_src)
        job_src = os.path.abspath(self.job.script)

        async def _deploy_one(node: Node):
            rc, _ = await node.conn.exec("test -f /workspace/worker.py")
            if rc != 0:
                await node.conn.push(worker_src, "/workspace/worker.py")
            else:
                log.debug(f"Node {node.rank}: worker.py already present, skipping push")
            await node.conn.push(job_src, "/workspace/job_module.py")

        results = await asyncio.gather(
            *[_deploy_one(n) for n in self._active_nodes()],
            return_exceptions=True,
        )
        for node, res in zip(self._active_nodes(), results):
            if isinstance(res, Exception):
                log.warning(f"Deploy failed on node {node.rank}: {res}")
                node.status = "dead"

        if not self._active_nodes():
            raise RuntimeError("Deploy failed on all nodes")

    async def run_round(
        self, state_dict: OrderedDict, round_num: int
    ) -> OrderedDict:
        """Single FedAvg round: push params -> workers train -> pull weights -> average."""
        round_t0 = time.monotonic()

        params_path = os.path.join(self.work_dir, "params.pt")
        torch.save(state_dict, params_path)

        active = self._active_nodes()
        if not active:
            raise RuntimeError("No active workers")

        if self.tracer:
            self.tracer.emit(
                "round_start",
                round_num=round_num,
                active_nodes=len(active),
            )

        # --- push params ---
        push_res = await asyncio.gather(
            *[n.conn.push(params_path, "/workspace/params.pt") for n in active],
            return_exceptions=True,
        )
        for node, res in zip(active, push_res):
            if isinstance(res, Exception):
                log.warning(f"Param push failed for node {node.rank}: {res}")
                node.status = "dead"

        active = self._active_nodes()
        if not active:
            raise RuntimeError("No workers alive after param push")

        # --- run training ---
        cmd_template = (
            "cd /workspace && python worker.py"
            " --params /workspace/params.pt"
            " --output /workspace/weights.pt"
            " --rank {rank}"
            f" --world-size {self.job.num_nodes}"
            f" --local-steps {self.job.local_steps}"
            f" --lr {self.job.lr}"
            f" --batch-size {self.job.batch_size}"
        )
        exec_res = await asyncio.gather(
            *[
                n.conn.exec(cmd_template.format(rank=n.rank), timeout=300.0)
                for n in active
            ],
            return_exceptions=True,
        )

        survivors = []
        for node, res in zip(active, exec_res):
            if isinstance(res, Exception):
                log.warning(f"Node {node.rank} raised: {res}")
                node.status = "dead"
            elif res[0] != 0:
                log.warning(f"Node {node.rank} exit {res[0]}: {res[1][:200]}")
                node.status = "dead"
            else:
                survivors.append(node)

        if not survivors:
            raise RuntimeError("All workers failed")

        # --- pull weights ---
        weight_dicts: list[OrderedDict] = []
        for node in survivors:
            local_path = os.path.join(self.work_dir, f"weights_{node.rank}.pt")
            try:
                await node.conn.pull("/workspace/weights.pt", local_path)
                w = torch.load(local_path, map_location="cpu", weights_only=True)
                weight_dicts.append(w)
            except Exception as e:
                log.warning(f"Weight pull failed for node {node.rank}: {e}")
                node.status = "dead"

        if not weight_dicts:
            raise RuntimeError("No weights collected")

        # --- aggregate ---
        agg_t0 = time.monotonic()
        new_state = average_weights(weight_dicts)
        agg_dur = time.monotonic() - agg_t0

        if self.tracer:
            weight_norm = sum(v.float().norm().item() for v in new_state.values())
            delta_norm = sum(
                (new_state[k].float() - state_dict[k].float()).norm().item()
                for k in new_state
            )
            self.tracer.emit(
                "aggregate",
                num_workers=len(weight_dicts),
                weight_norm=weight_norm,
                delta_norm=delta_norm,
                duration_s=agg_dur,
            )

        round_dur = time.monotonic() - round_t0

        if self.tracer:
            self.tracer.emit(
                "round_end",
                round_num=round_num,
                survivors=len(weight_dicts),
                total_nodes=len(self.nodes),
                duration_s=round_dur,
            )

        log.info(
            f"Round {round_num}: {len(weight_dicts)}/{len(self.nodes)} workers"
        )
        return new_state

    async def run(self) -> OrderedDict:
        """Full pipeline: provision -> deploy -> train rounds -> cleanup.

        Returns the final averaged state_dict.
        """
        job_mod = self._load_job_module()
        model = job_mod.make_model()
        state = model.state_dict()

        log.info(
            f"Job start: {self.job.rounds} rounds, "
            f"{self.job.num_nodes} nodes, "
            f"{self.job.local_steps} local steps/round"
        )

        await self.provision()
        await self.deploy()

        for r in range(self.job.rounds):
            state = await self.run_round(state, r)

        if self.tracer:
            self.tracer.summary()

        log.info("Job complete")
        await self.cleanup()
        return state

    async def cleanup(self):
        """Close connections, destroy instances."""
        from . import vastai

        for node in self.nodes:
            if node.conn:
                await node.conn.close()
            if node.vast_id:
                try:
                    await vastai.destroy_instance(node.vast_id)
                except Exception:
                    pass