import torch
import traceback
from core.device import DEVICE
from emotion.tracker import EmotionTracker
from reasoning.persona import PersonaController

class MVIState:
    def __init__(self):
        self.modules = {}
        self.health = {}

    def register(self, name, module, status=True):
        self.modules[name] = module
        self.health[name] = status

    def get(self, name):
        return self.modules.get(name)

    def report(self):
        return self.health


def safe_load(name, loader_fn, state):
    """
    Attempts to load a module and updates state health.
    """
    try:
        module = loader_fn()
        state.register(name, module, True)
        print(f"[LOAD SUCCESS] {name}")
    except Exception:
        traceback.print_exc()
        state.register(name, None, False)
        print(f"[LOAD FAILED] {name}")


def initialize_mvi():
    """
    Central MVI loader — loads all modules (language, emotion, vision, voice, memory)
    and registers them into MVIState with health tracking.
    """
    from language.tokenizer import SimpleTokenizer
    from language.embeddings import EmbeddingLayer
    from language.encoder import SentenceEncoder
    from language.intent import IntentClassifier
    from emotion.sentiment_model import SentimentRegressor
    from vision.image_encoder import ImageEncoder
    from vision.video_encoder import VideoEncoder
    from core.voice_encoder import VoiceEncoder
    from memory.short_term import ShortTermMemory
    from memory.long_term import LongTermMemory

    state = MVIState()

    # ---------- EMOTION TRACKER & PERSONA ----------
    state.register("emotion_tracker", EmotionTracker(window_size=6))
    state.register("persona", PersonaController())

    # ---------- LANGUAGE ----------
    def load_language():
        tokenizer = SimpleTokenizer()
        tokenizer.load_vocab("artifacts/vocab.json")

        embedder = EmbeddingLayer(
            tokenizer.vocab_size,
            pad_index=tokenizer.vocab[tokenizer.PAD_TOKEN]
        ).to(DEVICE)

        encoder = SentenceEncoder().to(DEVICE)
        encoder.load_state_dict(torch.load("artifacts/sentence_encoder.pt", map_location=DEVICE))

        intent = IntentClassifier(
            input_dim=encoder.projection.out_features,
            intent_labels=["question", "advice", "statement"]
        ).to(DEVICE)

        intent.load_state_dict(torch.load("artifacts/intent_classifier.pt", map_location=DEVICE))

        return {
            "tokenizer": tokenizer,
            "embedder": embedder.eval(),
            "encoder": encoder.eval(),
            "intent": intent.eval()
        }

    safe_load("language", load_language, state)

    # ---------- EMOTION ----------
    def load_emotion():
        model = SentimentRegressor(input_dim=128).to(DEVICE)
        model.load_state_dict(torch.load("artifacts/sentiment_regressor.pt", map_location=DEVICE))
        return model.eval()

    safe_load("emotion", load_emotion, state)

    # ---------- VISION ----------
    def load_image():
        model = ImageEncoder(embed_dim=128).to(DEVICE)
        model.load_state_dict(torch.load("artifacts/image_encoder.pt", map_location=DEVICE))
        return model.eval()

    safe_load("vision_image", load_image, state)

    def load_video():
        model = VideoEncoder(embed_dim=128).to(DEVICE)
        model.load_state_dict(torch.load("artifacts/video_encoder.pt", map_location=DEVICE))
        return model.eval()

    safe_load("vision_video", load_video, state)

    # ---------- VOICE ----------
    def load_voice():
        model = VoiceEncoder(embed_dim=128).to(DEVICE)
        model.load_state_dict(torch.load("artifacts/voice_encoder_best.pt", map_location=DEVICE))
        return model.eval()

    safe_load("voice", load_voice, state)

    # ---------- MEMORY ----------
    def load_memory():
        return {
            "stm": ShortTermMemory(max_len=10),
            "ltm": LongTermMemory(dim=128)
        }

    safe_load("memory", load_memory, state)

    print("\n[MVI SYSTEM HEALTH]")
    print(state.report())

    return state