"""Filesystem-based delta storage for the aggregation server."""

from __future__ import annotations

import json
import os
import re
import shutil
import time
from enum import Enum
from pathlib import Path

from chorus.exceptions import DuplicateClientError

import torch
from safetensors.torch import load_file, save_file


def _sanitize_client_id(client_id: str) -> str:
    """Sanitize client_id to prevent path traversal attacks."""
    # Strip path separators, null bytes, and leading dots
    safe = re.sub(r'[/\\:\x00]', '_', client_id)
    safe = safe.lstrip('.')
    # Cap length
    safe = safe[:128]
    if not safe:
        raise ValueError("client_id is empty or invalid after sanitization")
    return safe


class RoundState(str, Enum):
    OPEN = "open"
    AGGREGATING = "aggregating"
    CLOSED = "closed"


class DeltaStorage:
    """Stores LoRA deltas on the filesystem, organized by model and round."""

    def __init__(self, base_dir: str | Path = "./chorus_data"):
        self.base_dir = Path(base_dir)
        self.base_dir.mkdir(parents=True, exist_ok=True)

    def _model_dir(self, model_id: str) -> Path:
        safe_name = model_id.replace("/", "__")
        return self.base_dir / safe_name

    def _round_dir(self, model_id: str, round_id: int) -> Path:
        return self._model_dir(model_id) / f"round_{round_id}"

    def _deltas_dir(self, model_id: str, round_id: int) -> Path:
        return self._round_dir(model_id, round_id) / "deltas"

    def save_delta(
        self,
        model_id: str,
        round_id: int,
        client_id: str,
        tensors: dict,
        metadata: dict | None = None,
    ) -> Path:
        """Save a client's LoRA delta to disk."""
        client_id = _sanitize_client_id(client_id)
        deltas_dir = self._deltas_dir(model_id, round_id)
        deltas_dir.mkdir(parents=True, exist_ok=True)

        # Atomic duplicate detection via exclusive file creation
        lock_path = deltas_dir / f"{client_id}.lock"
        try:
            fd = os.open(str(lock_path), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
            os.close(fd)
        except FileExistsError:
            raise DuplicateClientError(
                f"Client '{client_id}' has already submitted to round {round_id}"
            )

        delta_path = deltas_dir / f"{client_id}.safetensors"
        save_file(tensors, str(delta_path))

        # Save metadata alongside
        meta_path = deltas_dir / f"{client_id}.json"
        meta = {
            "client_id": client_id,
            "model_id": model_id,
            "round_id": round_id,
            "timestamp": time.time(),
            **(metadata or {}),
        }
        meta_path.write_text(json.dumps(meta, indent=2))

        return delta_path

    def list_deltas(self, model_id: str, round_id: int) -> list[str]:
        """List client IDs that have submitted deltas for a given round."""
        deltas_dir = self._deltas_dir(model_id, round_id)
        if not deltas_dir.exists():
            return []
        return [p.stem for p in deltas_dir.glob("*.safetensors")]

    def load_delta(self, model_id: str, round_id: int, client_id: str) -> dict:
        """Load a single client's delta tensors."""
        delta_path = self._deltas_dir(model_id, round_id) / f"{client_id}.safetensors"
        return load_file(str(delta_path))

    def load_all_deltas(self, model_id: str, round_id: int) -> list[dict]:
        """Load all client deltas for a round."""
        client_ids = self.list_deltas(model_id, round_id)
        return [self.load_delta(model_id, round_id, cid) for cid in client_ids]

    def load_delta_metadata(self, model_id: str, round_id: int, client_id: str) -> dict:
        """Load metadata for a single client's delta."""
        meta_path = self._deltas_dir(model_id, round_id) / f"{client_id}.json"
        if not meta_path.exists():
            return {}
        return json.loads(meta_path.read_text())

    def load_all_delta_metadata(self, model_id: str, round_id: int) -> list[dict]:
        """Load metadata for all deltas in a round."""
        client_ids = self.list_deltas(model_id, round_id)
        return [self.load_delta_metadata(model_id, round_id, cid) for cid in client_ids]

    def save_aggregated(self, model_id: str, round_id: int, tensors: dict) -> Path:
        """Save the aggregated result for a round."""
        round_dir = self._round_dir(model_id, round_id)
        round_dir.mkdir(parents=True, exist_ok=True)

        agg_path = round_dir / "aggregated.safetensors"
        save_file(tensors, str(agg_path))

        # Update the latest symlink / pointer
        latest_path = self._model_dir(model_id) / "latest.safetensors"
        if latest_path.exists():
            latest_path.unlink()
        shutil.copy2(str(agg_path), str(latest_path))

        # Write round metadata
        meta_path = round_dir / "aggregation_meta.json"
        meta = {
            "model_id": model_id,
            "round_id": round_id,
            "num_clients": len(self.list_deltas(model_id, round_id)),
            "timestamp": time.time(),
        }
        meta_path.write_text(json.dumps(meta, indent=2))

        return agg_path

    def load_aggregated(self, model_id: str, round_id: int | None = None) -> dict | None:
        """Load aggregated tensors. If round_id is None, load latest."""
        if round_id is not None:
            agg_path = self._round_dir(model_id, round_id) / "aggregated.safetensors"
        else:
            agg_path = self._model_dir(model_id) / "latest.safetensors"

        if not agg_path.exists():
            return None
        return load_file(str(agg_path))

    def get_latest_round(self, model_id: str) -> int | None:
        """Get the latest completed round number for a model."""
        model_dir = self._model_dir(model_id)
        if not model_dir.exists():
            return None
        rounds = []
        for p in model_dir.iterdir():
            if p.is_dir() and p.name.startswith("round_"):
                agg = p / "aggregated.safetensors"
                if agg.exists():
                    rounds.append(int(p.name.split("_")[1]))
        return max(rounds) if rounds else None

    def get_current_round(self, model_id: str) -> int:
        """Get the current round (latest completed + 1, or 0 if none)."""
        latest = self.get_latest_round(model_id)
        return 0 if latest is None else latest + 1

    def cleanup_round(self, model_id: str, round_id: int) -> None:
        """Remove delta files for a completed round to save space."""
        deltas_dir = self._deltas_dir(model_id, round_id)
        if deltas_dir.exists():
            shutil.rmtree(deltas_dir)

    # --- Round state management ---

    def _round_state_path(self, model_id: str, round_id: int) -> Path:
        return self._round_dir(model_id, round_id) / "round_state.json"

    def get_round_state(self, model_id: str, round_id: int) -> RoundState:
        """Get the state of a round. Defaults to OPEN if no state file exists."""
        state_path = self._round_state_path(model_id, round_id)
        if not state_path.exists():
            return RoundState.OPEN
        data = json.loads(state_path.read_text())
        return RoundState(data["state"])

    def set_round_state(self, model_id: str, round_id: int, state: RoundState) -> None:
        """Set the state of a round."""
        round_dir = self._round_dir(model_id, round_id)
        round_dir.mkdir(parents=True, exist_ok=True)
        state_path = self._round_state_path(model_id, round_id)
        data = {"state": state.value, "timestamp": time.time()}
        state_path.write_text(json.dumps(data, indent=2))

    def is_round_accepting(self, model_id: str, round_id: int) -> bool:
        """Check if a round is still accepting submissions."""
        return self.get_round_state(model_id, round_id) == RoundState.OPEN

    # --- Base weight management ---

    def _base_weights_path(self, model_id: str) -> Path:
        return self._model_dir(model_id) / "base_weights.safetensors"

    def _base_weights_meta_path(self, model_id: str) -> Path:
        return self._model_dir(model_id) / "base_weights_meta.json"

    def save_base_weights(self, model_id: str, tensors: dict[str, torch.Tensor], meta: dict | None = None) -> Path:
        """Save base model weights for a model."""
        model_dir = self._model_dir(model_id)
        model_dir.mkdir(parents=True, exist_ok=True)
        path = self._base_weights_path(model_id)
        save_file(tensors, str(path))
        meta_path = self._base_weights_meta_path(model_id)
        meta_data = {"timestamp": time.time(), **(meta or {})}
        meta_path.write_text(json.dumps(meta_data, indent=2))
        return path

    def load_base_weights(self, model_id: str) -> dict[str, torch.Tensor] | None:
        """Load base model weights. Returns None if not set."""
        path = self._base_weights_path(model_id)
        if not path.exists():
            return None
        return load_file(str(path))

    def load_base_weights_meta(self, model_id: str) -> dict | None:
        """Load base weights metadata."""
        meta_path = self._base_weights_meta_path(model_id)
        if not meta_path.exists():
            return None
        return json.loads(meta_path.read_text())

    # --- Residual persistence ---

    def _residuals_path(self, model_id: str) -> Path:
        return self._model_dir(model_id) / "residuals.safetensors"

    def save_residuals(self, model_id: str, residuals: dict[str, torch.Tensor]) -> None:
        """Persist FedEx-LoRA residuals to disk."""
        if not residuals:
            # Remove residuals file if no residuals
            path = self._residuals_path(model_id)
            if path.exists():
                path.unlink()
            return
        model_dir = self._model_dir(model_id)
        model_dir.mkdir(parents=True, exist_ok=True)
        save_file(residuals, str(self._residuals_path(model_id)))

    def load_residuals(self, model_id: str) -> dict[str, torch.Tensor]:
        """Load persisted FedEx-LoRA residuals. Returns empty dict if none."""
        path = self._residuals_path(model_id)
        if not path.exists():
            return {}
        return load_file(str(path))

    def find_stuck_rounds(self, model_id: str) -> list[int]:
        """Find rounds stuck in AGGREGATING state (from a previous crash)."""
        model_dir = self._model_dir(model_id)
        if not model_dir.exists():
            return []
        stuck = []
        for p in model_dir.iterdir():
            if p.is_dir() and p.name.startswith("round_"):
                round_id = int(p.name.split("_")[1])
                if self.get_round_state(model_id, round_id) == RoundState.AGGREGATING:
                    stuck.append(round_id)
        return sorted(stuck)
