from __future__ import annotations

import os

from src.engines.sstgnn.engine import SSTGNNEngine
from src.services.media_utils import extract_video_frames


class SSTGNNModule:
    def __init__(self, cache_dir: str = "/data/model_cache"):
        os.environ.setdefault("MODEL_CACHE_DIR", cache_dir)
        self.engine = SSTGNNEngine()

    def score(self, video_path: str) -> dict:
        frames = extract_video_frames(video_path, max_frames=60)
        if not frames:
            return {"s3": 0.5, "vram_mb": 0}

        result = self.engine.run_video(frames)
        return {"s3": round(float(result.confidence), 4), "vram_mb": 0}