"""Model Manager Singleton for FastAPI ML Backend API (Milestone 7).

Loads BGE-small-en-v1.5, semantic_pca_v1.pkl, and dqn_triage_semantic_v1.zip ONCE at startup.
Executes deterministic 50D state building and Q-value inference without model mutation.
"""

import os
import sys
import time
import pickle
import numpy as np
import pandas as pd
from pathlib import Path
from typing import Dict, Any, Optional, List

import torch
from stable_baselines3 import DQN

# Ensure project root is in sys.path
PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from ml_rl.llm.semantic_encoder import PretrainedSemanticEncoder, build_semantic_text
from ml_rl.rl.state_builder import TriageStateBuilder, STATE_VERSION_SEMANTIC_V1

DQN_MODEL_PATH = PROJECT_ROOT / "experiments" / "models" / "dqn_triage_semantic_v1.zip"
PCA_MODEL_PATH = PROJECT_ROOT / "experiments" / "models" / "semantic_pca_v1.pkl"

PRIORITY_NAMES = ["Lowest", "Low", "Medium", "High", "Highest"]

FORBIDDEN_FIELDS = [
    "priority_score", "priority_name", "priority_id", "Priority", "ground_truth",
    "cycle_time", "reassignment_count", "reopen_count", "worklogs", "resolution"
]


class ModelManager:
    """Singleton model manager for loading and serving ML artifacts deterministically."""

    _instance: Optional["ModelManager"] = None

    def __init__(self):
        if ModelManager._instance is not None:
            raise RuntimeError("ModelManager is a singleton. Use ModelManager.get_instance().")

        print("Initializing ModelManager...")
        t0 = time.time()

        # Check required files
        if not DQN_MODEL_PATH.exists():
            raise FileNotFoundError(f"DQN Model Checkpoint not found at {DQN_MODEL_PATH}")
        if not PCA_MODEL_PATH.exists():
            raise FileNotFoundError(f"PCA Projection Model not found at {PCA_MODEL_PATH}")

        # 1. Load BGE Semantic Encoder
        self.encoder = PretrainedSemanticEncoder()
        
        # 2. Load PCA Model
        with open(PCA_MODEL_PATH, "rb") as f:
            self.pca = pickle.load(f)

        # 3. Load Frozen Semantic DQN
        self.dqn_model = DQN.load(DQN_MODEL_PATH)
        
        # 4. Initialize State Builder
        self.state_builder = TriageStateBuilder(state_version=STATE_VERSION_SEMANTIC_V1)

        self.device = self.encoder.device.upper()
        init_time = time.time() - t0
        print(f"ModelManager initialized on {self.device} in {round(init_time, 2)}s.")

    @classmethod
    def get_instance(cls) -> "ModelManager":
        if cls._instance is None:
            cls._instance = ModelManager()
        return cls._instance

    def predict(self, raw_input: Dict[str, Any]) -> Dict[str, Any]:
        """Performs deterministic 50D inference on pre-triage issue input dict."""
        t_start = time.time()

        # Anti-Leakage Check: Strip or reject forbidden target & post-triage outcome fields
        clean_input = {k: v for k, v in raw_input.items() if k not in FORBIDDEN_FIELDS}

        # 1. Build Triage-Safe Text Prompt
        semantic_text = build_semantic_text(clean_input)

        # 2. Extract 384D BGE Embedding
        emb_384 = self.encoder.encode_single_text(semantic_text)  # (384,) float32

        # 3. Transform via PCA -> 29D
        emb_29 = self.pca.transform(emb_384.reshape(1, -1))[0].astype(np.float32)  # (29,) float32

        # 4. Build 50D Observation Vector (21D Base + 29D Semantic)
        obs_50 = self.state_builder.build_state(clean_input, llm_features=emb_29)  # (50,) float32

        # 5. Execute Deterministic Inference with Q-values
        obs_tensor = torch.as_tensor(obs_50, dtype=torch.float32).unsqueeze(0).to(self.dqn_model.device)
        with torch.no_grad():
            q_values_arr = self.dqn_model.q_net(obs_tensor).cpu().numpy()[0]

        action_label = int(np.argmax(q_values_arr))
        action_name = PRIORITY_NAMES[action_label]

        latency_ms = round((time.time() - t_start) * 1000.0, 2)

        q_dict = {
            name: round(float(val), 4) for name, val in zip(PRIORITY_NAMES, q_values_arr)
        }

        return {
            "success": True,
            "prediction": {
                "label": action_label,
                "name": action_name
            },
            "model": {
                "name": "Semantic DQN",
                "checkpoint": "dqn_triage_semantic_v1",
                "state_version": "triage_semantic_v1",
                "state_dimension": 50
            },
            "semantic": {
                "model": "BAAI/bge-small-en-v1.5",
                "embedding_dimension": 384,
                "pca_dimension": 29
            },
            "q_values": q_dict,
            "inference": {
                "device": self.device,
                "latency_ms": latency_ms
            }
        }

    def predict_batch(self, raw_issues: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Performs batch 50D inference and ranks issues descending by expected priority score."""
        if not raw_issues:
            raise ValueError("Backlog request cannot be empty.")
        if len(raw_issues) > 20:
            raise ValueError("Maximum backlog size exceeded (limit: 20 issues).")

        t_start = time.time()
        results = []

        for idx, raw_input in enumerate(raw_issues):
            # Anti-Leakage Check: Strip forbidden target & post-triage outcome fields
            clean_input = {k: v for k, v in raw_input.items() if k not in FORBIDDEN_FIELDS}

            # Map summary fields
            summary_text = clean_input.get("summary") or clean_input.get("issue_summary") or clean_input.get("title") or ""
            clean_input["issue_summary"] = summary_text

            # Map severity/type fields if present under standard names
            if "issue_type" in clean_input:
                clean_input["issuetype_name"] = clean_input["issue_type"]
            if "severity" in clean_input:
                clean_input["severity_value"] = clean_input["severity"]

            # Map boolean & numeric flags
            clean_input["Attachment Count"] = clean_input.get("attachment_count", 0)
            clean_input["Component Count"] = clean_input.get("component_count", 0)
            clean_input["Fix Version Count"] = clean_input.get("fix_version_count", 0)
            clean_input["Label Count"] = clean_input.get("label_count", 0)
            clean_input["Has Epic Link"] = clean_input.get("has_epic_link", False)

            # 1. Build Triage-Safe Text Prompt
            semantic_text = build_semantic_text(clean_input)

            # 2. Extract 384D BGE Embedding
            emb_384 = self.encoder.encode_single_text(semantic_text)

            # 3. Transform via PCA -> 29D
            emb_29 = self.pca.transform(emb_384.reshape(1, -1))[0].astype(np.float32)

            # 4. Build 50D Observation Vector
            obs_50 = self.state_builder.build_state(clean_input, llm_features=emb_29)

            # 5. Execute Deterministic Inference with Q-values
            obs_tensor = torch.as_tensor(obs_50, dtype=torch.float32).unsqueeze(0).to(self.dqn_model.device)
            with torch.no_grad():
                q_values_arr = self.dqn_model.q_net(obs_tensor).cpu().numpy()[0]

            action_label = int(np.argmax(q_values_arr))
            action_name = PRIORITY_NAMES[action_label]

            # Softmax calculation for expected priority score: sum_a [ a * softmax(Q(s,a)) ]
            exp_q = np.exp(q_values_arr - np.max(q_values_arr))
            softmax_probs = exp_q / np.sum(exp_q)
            expected_score = float(np.sum(np.arange(5) * softmax_probs))

            q_dict = {
                name: round(float(val), 4) for name, val in zip(PRIORITY_NAMES, q_values_arr)
            }

            results.append({
                "original_index": idx,
                "summary": summary_text,
                "description": clean_input.get("description", ""),
                "issue_type": clean_input.get("issue_type", "Bug"),
                "severity": clean_input.get("severity", "Normal"),
                "priority": action_name,
                "action": action_label,
                "expected_priority_score": round(expected_score, 4),
                "q_values": q_dict,
                "state_50d_preview": [round(float(v), 4) for v in obs_50[:10]]
            })

        # Rank descending by expected_priority_score
        results.sort(key=lambda x: x["expected_priority_score"], reverse=True)

        for rank_idx, item in enumerate(results, start=1):
            item["rank"] = rank_idx

        latency_ms = round((time.time() - t_start) * 1000.0, 2)

        return {
            "success": True,
            "count": len(results),
            "results": results,
            "device": self.device,
            "latency_ms": latency_ms
        }

    def simulate_single_task(self, raw_input: Dict[str, Any]) -> Dict[str, Any]:
        """Performs task-level 50D state construction and Q-value evaluation for live task simulation."""
        t_start = time.time()

        # Anti-Leakage Check: Strip forbidden target & post-triage outcome fields
        clean_input = {k: v for k, v in raw_input.items() if k not in FORBIDDEN_FIELDS}

        summary_text = clean_input.get("summary") or clean_input.get("issue_summary") or clean_input.get("title") or ""
        clean_input["issue_summary"] = summary_text

        if "issue_type" in clean_input:
            clean_input["issuetype_name"] = clean_input["issue_type"]
        if "severity" in clean_input:
            clean_input["severity_value"] = clean_input["severity"]

        clean_input["Attachment Count"] = clean_input.get("attachment_count", 0)
        clean_input["Component Count"] = clean_input.get("component_count", 0)
        clean_input["Fix Version Count"] = clean_input.get("fix_version_count", 0)
        clean_input["Label Count"] = clean_input.get("label_count", 0)
        clean_input["Has Epic Link"] = clean_input.get("has_epic_link", False)

        # 1. Build Triage-Safe Text Prompt
        semantic_text = build_semantic_text(clean_input)

        # 2. Extract 384D BGE Embedding
        emb_384 = self.encoder.encode_single_text(semantic_text)

        # 3. Transform via PCA -> 29D
        emb_29 = self.pca.transform(emb_384.reshape(1, -1))[0].astype(np.float32)

        # 4. Build 50D Observation Vector
        obs_50 = self.state_builder.build_state(clean_input, llm_features=emb_29)

        # 5. Execute Deterministic Inference with Q-values
        obs_tensor = torch.as_tensor(obs_50, dtype=torch.float32).unsqueeze(0).to(self.dqn_model.device)
        with torch.no_grad():
            q_values_arr = self.dqn_model.q_net(obs_tensor).cpu().numpy()[0]

        action_label = int(np.argmax(q_values_arr))
        action_name = PRIORITY_NAMES[action_label]

        # Softmax calculation for expected priority score
        exp_q = np.exp(q_values_arr - np.max(q_values_arr))
        softmax_probs = exp_q / np.sum(exp_q)
        expected_score = float(np.sum(np.arange(5) * softmax_probs))

        q_dict = {
            name: round(float(val), 4) for name, val in zip(PRIORITY_NAMES, q_values_arr)
        }

        latency_ms = round((time.time() - t_start) * 1000.0, 2)

        return {
            "success": True,
            "action": action_label,
            "priority": action_name,
            "q_values": q_dict,
            "expected_priority_score": round(expected_score, 4),
            "state_dimension": 50,
            "semantic_dimension": 29,
            "state_50d_preview": [round(float(v), 4) for v in obs_50[:10]],
            "device": self.device,
            "latency_ms": latency_ms
        }