from pathlib import Path
from os import getenv

from score_miner_core.detector.detector_router import create_detector
from score_miner_core.runtime.miner_runtime import MinerRuntime
from score_miner_core.runtime.postprocess import PostprocessConfig
from score_miner_core.runtime.role_cleanup import RoleCleanupConfig
from score_miner_core.runtime.team_color import TeamColorConfig
from score_miner_core.runtime.tracking import TrackingConfig


class Miner:
    def __init__(self, path_hf_repo: Path) -> None:
        detector_name = getenv("SCORE_MINER_DETECTOR", "rfdetr_m")
        threshold = float(getenv("SCORE_MINER_THRESHOLD", "0.75"))
        player_cls_id = int(getenv("SCORE_MINER_PLAYER_CLS_ID", "0"))
        ball_cls_id_raw = getenv("SCORE_MINER_BALL_CLS_ID")
        ball_cls_id = int(ball_cls_id_raw) if ball_cls_id_raw else None
        input_color_space = getenv("SCORE_MINER_INPUT_COLOR_SPACE", "bgr")
        postprocess_config = PostprocessConfig.from_env()
        team_color_config = TeamColorConfig.from_env()
        tracking_config = TrackingConfig.from_env()
        role_cleanup_config = RoleCleanupConfig.from_env()
        detector = create_detector(
            detector_name,
            threshold=threshold,
            player_cls_id=player_cls_id,
            ball_cls_id=ball_cls_id,
        )
        self.runtime = MinerRuntime(
            path_hf_repo,
            detector=detector,
            postprocess_config=postprocess_config,
            team_color_config=team_color_config,
            tracking_config=tracking_config,
            role_cleanup_config=role_cleanup_config,
            input_color_space=input_color_space,
        )

    def __repr__(self) -> str:
        return repr(self.runtime)

    def predict_batch(self, batch_images, offset: int, n_keypoints: int):
        return self.runtime.predict_batch(batch_images, offset, n_keypoints)