import os
import torch
import logging

logger = logging.getLogger(__name__)

from model import RouterNetwork


class RouterPolicy:
    """Thin wrapper around RouterNetwork for loading, inference and checkpointing."""

    def __init__(self, model: RouterNetwork, device: str = "cpu"):
        self.model = model.to(device)
        self.device = torch.device(device)
        self.model.eval()

    @classmethod
    def load(cls, checkpoint_path: str, embedding_dim: int | None = None, device: str = "cpu"):
        """Load RouterNetwork from checkpoint. If checkpoint contains container dict with
        'model_state_dict' will use that, otherwise will try to load state_dict directly.

        embedding_dim may be required if checkpoint does not include meta info.
        """
        checkpoint_path = os.path.abspath(checkpoint_path)
        if not os.path.exists(checkpoint_path):
            raise FileNotFoundError(f"Router checkpoint not found: {checkpoint_path}")

        sd = torch.load(checkpoint_path, map_location="cpu")

        # Unpack common formats
        if isinstance(sd, dict) and "model_state_dict" in sd:
            model_state = sd["model_state_dict"]
            meta = sd.get("meta", {})
        else:
            model_state = sd
            meta = {}

        # Infer embedding dim from meta if available
        emb_dim = embedding_dim or meta.get("embedding_dim") or meta.get("embed_dim") or 384

        model = RouterNetwork(embedding_dim=emb_dim)

        missing, unexpected = model.load_state_dict(model_state, strict=False)
        if missing or unexpected:
            logger.info(f"RouterPolicy.load: missing keys={list(missing)} unexpected={list(unexpected)}")

        return cls(model, device=device)

    def predict_sync(self, emb_tensor: torch.Tensor):
        """Synchronous prediction: accepts a 1-D or 2-D embedding tensor and returns
        (action_idx:int, prob:float, value:float).
        """
        was_1d = False
        if emb_tensor.dim() == 1:
            emb_tensor = emb_tensor.unsqueeze(0)
            was_1d = True

        emb = emb_tensor.to(self.device)

        with torch.no_grad():
            logits, value = self.model(emb)
            probs = torch.softmax(logits, dim=-1)
            top_idx = int(probs.argmax(dim=-1)[0].item())
            top_prob = float(probs[0, top_idx].item())
            val = float(value.squeeze(1)[0].item()) if value is not None else 0.0

        return top_idx, top_prob, val

    def save_checkpoint(self, path: str, optimizer: torch.optim.Optimizer | None = None, meta: dict | None = None):
        path = os.path.abspath(path)
        payload = {
            "model_state_dict": self.model.state_dict(),
            "meta": meta or {}
        }
        if optimizer is not None:
            try:
                payload["optimizer_state_dict"] = optimizer.state_dict()
            except Exception:
                logger.warning("Could not include optimizer state in checkpoint")

        torch.save(payload, path)
        return path