"""
Commander AI Lab — Policy Inference Server
════════════════════════════════════════════
FastAPI service that loads a trained PolicyNetwork checkpoint and serves
macro-action predictions via HTTP.

The Java batch runner calls this endpoint at each decision point
to get the learned policy's action instead of using Forge's built-in AI.

Endpoints:
    POST /api/ml/predict   — Predict action from raw game state JSON
    GET  /api/ml/model     — Current model info (loaded checkpoint, accuracy)
    POST /api/ml/reload    — Hot-reload a new checkpoint without restart

Design:
    - Loads card embeddings + state encoder on startup
    - Accepts DecisionSnapshot-shaped JSON from Java
    - Encodes state → tensor, runs model forward pass, returns action
    - Supports temperature-controlled sampling or greedy argmax
    - Thread-safe via torch.no_grad() inference

Usage (standalone):
    python -m ml.serving.policy_server --port 8090
    python -m ml.serving.policy_server --checkpoint ml/models/checkpoints/best_policy.pt

Usage (integrated into lab_api.py):
    Import PolicyInferenceService and register endpoints on the existing app.
"""

import logging
import os
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional

import numpy as np

project_root = str(Path(__file__).parent.parent.parent)
if project_root not in sys.path:
    sys.path.insert(0, project_root)

# Conditional imports
try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

from ml.config.scope import (
    NUM_ACTIONS, IDX_TO_ACTION, MacroAction, STATE_DIMS,
)
from ml.encoder.state_encoder import CardEmbeddingIndex, StateEncoder

# Unique logger name so log lines are attributable to this module specifically
logger = logging.getLogger("ml.serving.policy_server")


def _unwrap_logits(raw):
    """
    Normalise whatever self.model() returns to a plain logits tensor.

    PolicyNetwork.forward() returns a raw tensor, but the model may be
    wrapped (DataParallel, torch.compile) or accidentally swapped for
    PolicyValueNetwork whose forward_with_value() returns (logits, value).
    A dict return would cause 'dict / float' TypeError on the temperature
    division -- guard against all three cases here.
    """
    if isinstance(raw, (tuple, list)):
        # (logits, value) from PolicyValueNetwork or similar
        return raw[0]
    if isinstance(raw, dict):
        # e.g. {"logits": tensor, "value": tensor}
        for key in ("logits", "policy", "action_logits"):
            if key in raw:
                return raw[key]
        # Fallback: grab the first tensor value in the dict
        for v in raw.values():
            if isinstance(v, torch.Tensor):
                return v
        raise TypeError(
            f"model() returned a dict with no recognisable logits key: {list(raw.keys())}"
        )
    # Plain tensor -- the normal case
    return raw


class PolicyInferenceService:
    """
    Loads a trained PolicyNetwork and provides action predictions.

    Thread-safe -- uses torch.no_grad() for inference.
    Designed to be embedded in lab_api.py or run standalone.
    """

    def __init__(
        self,
        checkpoint_dir: str = None,
        embeddings_dir: str = None,
        device: str = None,
    ):
        self.checkpoint_dir = checkpoint_dir or os.path.join(
            project_root, "ml", "models", "checkpoints"
        )
        self.embeddings_dir = embeddings_dir or os.path.join(
            project_root, "embeddings"
        )
        self.device = device or self._detect_device()

        self.model = None
        self.encoder = None
        self.card_index = None
        self.checkpoint_info: Dict = {}
        self._loaded = False
        self._load_error: Optional[str] = None

    def _detect_device(self) -> str:
        if not TORCH_AVAILABLE:
            return "cpu"
        if torch.cuda.is_available():
            return "cuda"
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
        return "cpu"

    def load(self) -> bool:
        """
        Load card embeddings, state encoder, and policy model.
        Returns True if everything loaded successfully.
        """
        if not TORCH_AVAILABLE:
            self._load_error = "PyTorch not installed"
            logger.error("PyTorch not installed -- cannot serve policy")
            return False

        # 1. Load card embeddings
        self.card_index = CardEmbeddingIndex(self.embeddings_dir)
        if not self.card_index.load():
            self._load_error = "Card embeddings not found"
            logger.error(
                "Card embeddings not found at %s. "
                "Start the lab server to auto-download, or run coach first.",
                self.embeddings_dir,
            )
            return False

        self.encoder = StateEncoder(self.card_index)

        # 2. Load policy checkpoint
        checkpoint_path = os.path.join(self.checkpoint_dir, "best_policy.pt")
        if not os.path.exists(checkpoint_path):
            self._load_error = f"No checkpoint at {checkpoint_path}"
            logger.warning("No trained model found at %s", checkpoint_path)
            logger.warning("Train a model first: python -m ml.scripts.ml_cli train")
            return False

        try:
            from ml.training.policy_network import load_checkpoint
            self.model, checkpoint = load_checkpoint(checkpoint_path, self.device)
            self.model.eval()

            self.checkpoint_info = {
                "path": checkpoint_path,
                "epoch": checkpoint.get("epoch", "?"),
                "metrics": checkpoint.get("metrics", {}),
                "device": self.device,
            }

            self._loaded = True
            self._load_error = None
            logger.info(
                "Policy model loaded (epoch %s, device=%s)",
                self.checkpoint_info["epoch"], self.device,
            )
            return True

        except Exception as e:
            self._load_error = str(e)
            logger.error("Failed to load checkpoint: %s", e, exc_info=True)
            return False

    def reload(self, checkpoint_path: str = None) -> bool:
        """Hot-reload a checkpoint without reloading embeddings."""
        if not TORCH_AVAILABLE:
            return False

        path = checkpoint_path or os.path.join(self.checkpoint_dir, "best_policy.pt")
        if not os.path.exists(path):
            self._load_error = f"Checkpoint not found: {path}"
            return False

        try:
            from ml.training.policy_network import load_checkpoint
            self.model, checkpoint = load_checkpoint(path, self.device)
            self.model.eval()

            self.checkpoint_info = {
                "path": path,
                "epoch": checkpoint.get("epoch", "?"),
                "metrics": checkpoint.get("metrics", {}),
                "device": self.device,
            }
            self._loaded = True
            self._load_error = None
            logger.info("Model hot-reloaded from %s (epoch %s)",
                        path, self.checkpoint_info["epoch"])
            return True

        except Exception as e:
            self._load_error = str(e)
            logger.error("Hot-reload failed: %s", e, exc_info=True)
            return False

    def predict(
        self,
        decision_snapshot: Dict,
        playstyle: str = "midrange",
        temperature: float = 1.0,
        greedy: bool = False,
    ) -> Dict:
        """
        Predict a macro-action from a raw decision snapshot.

        Args:
            decision_snapshot: DecisionSnapshot JSON from Java
                Must contain: turn, phase, active_seat, players[]
            playstyle: Deck playstyle hint ("aggro", "control", "midrange", "combo")
            temperature: Softmax temperature (lower = more deterministic)
            greedy: If True, always pick argmax (ignore temperature)

        Returns:
            {
                "action": "cast_creature",
                "action_index": 0,
                "confidence": 0.73,
                "probabilities": {"cast_creature": 0.73, "pass": 0.12, ...},
                "inference_ms": 2.3,
            }
        """
        if not self._loaded:
            return {
                "error": "Model not loaded",
                "detail": self._load_error or "Call load() first",
            }

        t_start = time.time()

        try:
            # Encode the raw snapshot to a state vector
            state_vec = self.encoder.encode(decision_snapshot, playstyle)
            state_tensor = torch.from_numpy(
                state_vec.astype(np.float32)
            ).unsqueeze(0).to(self.device)

            # Run inference
            with torch.no_grad():
                raw_output = self.model(state_tensor)
                # Guard: unwrap tuple/dict returns (DataParallel, torch.compile,
                # or accidental PolicyValueNetwork swap) before any arithmetic.
                logits = _unwrap_logits(raw_output)

                if greedy:
                    action_idx = logits.argmax(dim=-1).item()
                    probs = torch.softmax(logits, dim=-1).cpu().numpy().flatten()
                else:
                    scaled = logits / max(temperature, 0.01)
                    probs = torch.softmax(scaled, dim=-1).cpu().numpy().flatten()
                    action_idx = int(np.random.choice(NUM_ACTIONS, p=probs))

            elapsed_ms = (time.time() - t_start) * 1000

            action = IDX_TO_ACTION[action_idx]
            confidence = float(probs[action_idx])

            # Build probability map
            prob_map = {}
            for i in range(NUM_ACTIONS):
                prob_map[IDX_TO_ACTION[i].value] = round(float(probs[i]), 4)

            return {
                "action": action.value,
                "action_index": action_idx,
                "confidence": round(confidence, 4),
                "probabilities": prob_map,
                "inference_ms": round(elapsed_ms, 2),
            }

        except Exception as e:
            elapsed_ms = (time.time() - t_start) * 1000
            logger.error("Inference error: %s", e, exc_info=True)
            return {
                "error": str(e),
                "inference_ms": round(elapsed_ms, 2),
            }

    def predict_batch(
        self,
        snapshots: List[Dict],
        playstyle: str = "midrange",
        greedy: bool = True,
    ) -> List[Dict]:
        """Predict actions for multiple snapshots at once (vectorized)."""
        if not self._loaded:
            return [{"error": "Model not loaded"}] * len(snapshots)

        t_start = time.time()
        results = []

        try:
            # Encode all states
            state_vecs = []
            for snap in snapshots:
                ps = snap.get("archetype", playstyle)
                vec = self.encoder.encode(snap, ps)
                state_vecs.append(vec)

            states = np.stack(state_vecs).astype(np.float32)
            state_tensor = torch.from_numpy(states).to(self.device)

            # Batch inference
            with torch.no_grad():
                raw_output = self.model(state_tensor)
                # Same guard as predict() -- unwrap before softmax/argmax.
                logits = _unwrap_logits(raw_output)
                probs = torch.softmax(logits, dim=-1).cpu().numpy()

                if greedy:
                    action_indices = logits.argmax(dim=-1).cpu().numpy()
                else:
                    action_indices = np.array([
                        np.random.choice(NUM_ACTIONS, p=p) for p in probs
                    ])

            elapsed_ms = (time.time() - t_start) * 1000

            for i, (idx, prob) in enumerate(zip(action_indices, probs)):
                action = IDX_TO_ACTION[int(idx)]
                results.append({
                    "action": action.value,
                    "action_index": int(idx),
                    "confidence": round(float(prob[idx]), 4),
                })

            logger.debug(
                "Batch inference: %d samples in %.1fms (%.2f ms/sample)",
                len(snapshots), elapsed_ms, elapsed_ms / max(len(snapshots), 1),
            )

        except Exception as e:
            logger.error("Batch inference error: %s", e, exc_info=True)
            results = [{"error": str(e)}] * len(snapshots)

        return results

    def get_status(self) -> Dict:
        """Return current service status."""
        return {
            "loaded": self._loaded,
            "error": self._load_error,
            "device": self.device,
            "torch_available": TORCH_AVAILABLE,
            "checkpoint": self.checkpoint_info if self._loaded else None,
            "embeddings_loaded": (
                self.card_index is not None and self.card_index._loaded
            ),
        }


# ==========================================================
# Standalone FastAPI Server
# ==========================================================

def create_standalone_app():
    """Create a standalone FastAPI app for the policy server."""
    from fastapi import FastAPI, HTTPException
    from pydantic import BaseModel

    app = FastAPI(
        title="Commander AI Lab -- Policy Server",
        description="Serves macro-action predictions from a trained policy network",
        version="1.0.0",
    )

    service = PolicyInferenceService()

    class PredictRequest(BaseModel):
        """Request body for /predict endpoint."""
        # Game state (matching DecisionSnapshot from Java)
        turn: int = 1
        phase: str = "main_1"
        active_seat: int = 0
        players: List[Dict] = []
        action: Optional[Dict] = None

        # Optional metadata
        game_id: str = ""
        archetype: str = "midrange"
        commander: str = ""
        deck_name: str = ""

        # Inference parameters
        temperature: float = 1.0
        greedy: bool = False

    class PredictResponse(BaseModel):
        action: str
        action_index: int
        confidence: float
        probabilities: Dict[str, float] = {}
        inference_ms: float = 0.0
        error: Optional[str] = None

    @app.on_event("startup")
    async def startup():
        service.load()

    @app.post("/api/ml/predict", response_model=PredictResponse)
    async def predict(req: PredictRequest):
        if not service._loaded:
            raise HTTPException(503, detail=service._load_error or "Model not loaded")

        snapshot = req.dict(exclude={"temperature", "greedy"})
        result = service.predict(
            snapshot,
            playstyle=req.archetype,
            temperature=req.temperature,
            greedy=req.greedy,
        )

        if "error" in result:
            raise HTTPException(500, detail=result["error"])
        return result

    @app.get("/api/ml/model")
    async def model_info():
        return service.get_status()

    @app.post("/api/ml/reload")
    async def reload_model(checkpoint: str = None):
        ok = service.reload(checkpoint)
        return {
            "success": ok,
            "status": service.get_status(),
        }

    @app.get("/health")
    async def health():
        return {
            "status": "ok" if service._loaded else "degraded",
            "model_loaded": service._loaded,
        }

    return app, service


def main():
    """Run standalone policy server."""
    import argparse
    import uvicorn

    parser = argparse.ArgumentParser(
        description="Commander AI Lab -- Policy Inference Server"
    )
    parser.add_argument("--port", type=int, default=8090)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--checkpoint-dir", default=None)
    parser.add_argument("--embeddings-dir", default=None)
    parser.add_argument("--device", default=None)
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="[%(name)s] %(message)s")

    # Override defaults if provided
    if args.checkpoint_dir or args.embeddings_dir or args.device:
        # Re-create the service will be handled at startup
        pass

    app, _ = create_standalone_app()
    uvicorn.run(app, host=args.host, port=args.port)


if __name__ == "__main__":
    main()