api/deps.py

"""
ArcisVLM API — Dependency injection for model, agents, and camera manager.

Provides a singleton ModelManager that lazy-loads the VLM and agent framework
on first use, and exposes it via get_model_manager().
"""

from __future__ import annotations

import logging
import os
import time
import threading
from typing import Optional

logger = logging.getLogger(__name__)

# Singleton instance
_manager: Optional["ModelManager"] = None
_lock = threading.Lock()


class ModelManager:
    """
    Manages the lifecycle of the VLM model, agent orchestrator, and camera manager.

    All heavy objects are lazy-loaded on initialize() to keep import time fast.
    """

    def __init__(
        self,
        config_path: str = "configs/scale_1.3b.yaml",
        checkpoint_path: Optional[str] = None,
        device: str = "cpu",
    ) -> None:
        self.config_path = config_path
        self.checkpoint_path = checkpoint_path
        self.device = device

        self._model = None
        self._tokenizer = None
        self._orchestrator = None
        self._camera_manager = None
        self._initialized = False
        self._alert_history: list[dict] = []
        self.latency_history: list[float] = []
        self.inference_count: int = 0
        self.queries_per_sec: float = 0.0
        self.start_time: float = 0.0
        self.model_params: int = 0
        self._gemma_backbone = None  # Gemma 4 E2B backbone (when available)

    @property
    def is_initialized(self) -> bool:
        return self._initialized

    @property
    def model(self):
        return self._model

    @property
    def tokenizer(self):
        return self._tokenizer

    @property
    def orchestrator(self):
        return self._orchestrator

    @property
    def camera_manager(self):
        return self._camera_manager

    @property
    def gemma_backbone(self):
        return self._gemma_backbone

    @property
    def alert_history(self) -> list[dict]:
        return self._alert_history

    def add_alert_event(self, event: dict) -> None:
        self._alert_history.append(event)
        # Keep only last 1000 events
        if len(self._alert_history) > 1000:
            self._alert_history = self._alert_history[-1000:]

    def initialize(self) -> None:
        """Load model, tokenizer, and set up the agent orchestrator."""
        if self._initialized:
            return

        logger.info("ModelManager: initializing (config=%s, device=%s)", self.config_path, self.device)

        try:
            import yaml
            import torch

            # Load config
            with open(self.config_path) as f:
                config = yaml.safe_load(f)

            # Build model
            from model.vlm import VLJEPAModel
            self._model = VLJEPAModel(config)

            # Load checkpoint if provided
            if self.checkpoint_path:
                ckpt = torch.load(self.checkpoint_path, map_location=self.device, weights_only=False)
                state = ckpt.get("model_state_dict", ckpt)
                self._model.load_state_dict(state, strict=False)
                logger.info("Loaded checkpoint: %s", self.checkpoint_path)

            self._model = self._model.to(self.device)
            self._model.eval()

            # Load tokenizer
            from model.tokenizer import BPETokenizer
            tok_path = config.get("tokenizer", {}).get("path", "checkpoints/tokenizer_32k.json")
            self._tokenizer = BPETokenizer(vocab_size=config.get("tokenizer", {}).get("vocab_size", 32768))
            try:
                self._tokenizer.load(tok_path)
                logger.info("Loaded tokenizer from %s", tok_path)
            except FileNotFoundError:
                logger.warning("Tokenizer file not found at %s — using untrained tokenizer", tok_path)

            # Set up orchestrator with agents
            self._setup_agents()

            self.start_time = time.time()
            if self._model:
                self.model_params = sum(p.numel() for p in self._model.parameters())

            # Try loading Gemma 4 E2B backbone (if USE_GEMMA=1 env var set)
            if os.environ.get("USE_GEMMA", "").strip() == "1":
                try:
                    from model.gemma_backbone import GemmaBackbone
                    self._gemma_backbone = GemmaBackbone(device=self.device)
                    self._gemma_backbone.load()
                    self.model_params = sum(
                        p.numel() for p in self._gemma_backbone._model.parameters()
                    )
                    logger.info("Gemma 4 E2B backbone loaded successfully")
                except Exception as e:
                    logger.warning(f"Gemma backbone not available: {e}. Using custom VL-JEPA.")
                    self._gemma_backbone = None

            self._initialized = True
            logger.info("ModelManager: initialization complete")

        except Exception as e:
            logger.error("ModelManager: initialization failed — %s", e)
            # Still mark as partially initialized so API can respond with degraded status
            self._initialized = True
            logger.info("ModelManager: running in degraded mode (no model)")

    def _setup_agents(self) -> None:
        """Create the MotherOrchestrator and register child agents."""
        from agents.mother import MotherOrchestrator
        from agents.vqa import VQAAgent
        from agents.detect import DetectAgent
        from agents.alert import AlertAgent
        from agents.caption import CaptionAgent
        from agents.track import TrackingAgent
        from agents.count import CountingAgent
        from agents.ocr import OCRAgent
        from agents.reason import ReasoningAgent

        self._orchestrator = MotherOrchestrator(max_agents_per_type=4)

        # Create all 8 agent types backed by the shared model
        agent_classes = [
            ("vqa-0001", "vqa", VQAAgent),
            ("detect-0001", "detect", DetectAgent),
            ("alert-0001", "alert", AlertAgent),
            ("caption-0001", "caption", CaptionAgent),
            ("track-0001", "track", TrackingAgent),
            ("count-0001", "count", CountingAgent),
            ("ocr-0001", "ocr", OCRAgent),
            ("reason-0001", "reason", ReasoningAgent),
        ]

        for agent_id, expert_type, cls in agent_classes:
            try:
                agent = cls(agent_id=agent_id, model=self._model, tokenizer=self._tokenizer)
                self._orchestrator.register_agent(agent)
            except Exception as e:
                logger.warning("Failed to register %s agent: %s", expert_type, e)

    def submit_task(self, task):
        """Submit a task to the orchestrator (synchronous path)."""
        if self._orchestrator is None:
            from agents.base import Result
            return Result(answer="", error="Agent orchestrator not initialized")
        start = time.time()
        result = self._orchestrator.submit(task)
        elapsed_ms = (time.time() - start) * 1000
        self.latency_history.append(elapsed_ms)
        self.inference_count += 1
        if len(self.latency_history) > 10000:
            self.latency_history = self.latency_history[-10000:]
        return result

    def shutdown(self) -> None:
        """Clean up all resources."""
        if self._orchestrator:
            self._orchestrator.shutdown()
        if self._camera_manager:
            try:
                self._camera_manager.stop_all()
            except Exception:
                pass
        self._initialized = False
        logger.info("ModelManager: shutdown complete")


def get_model_manager() -> ModelManager:
    """Return the global ModelManager singleton, creating it if needed."""
    global _manager
    if _manager is None:
        with _lock:
            if _manager is None:
                _manager = ModelManager()
    return _manager


def configure_model_manager(
    config_path: str = "configs/scale_1.3b.yaml",
    checkpoint_path: Optional[str] = None,
    device: str = "cpu",
) -> ModelManager:
    """Create and set the global ModelManager with specific options."""
    global _manager
    with _lock:
        _manager = ModelManager(
            config_path=config_path,
            checkpoint_path=checkpoint_path,
            device=device,
        )
    return _manager