"""The project's CNN + Transformer as a :class:`Captioner`."""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Any, Protocol

from captioning.baselines.base import Captioner, CaptionerIdentity
from captioning.config.schema import AppConfig, ComparedModelConfig

if TYPE_CHECKING:
    from captioning.inference.predictor import DecodeStrategy


class _ModelSettings(Protocol):
    max_length: int


class _PredictorConfig(Protocol):
    @property
    def model(self) -> _ModelSettings: ...


class PathPredictor(Protocol):
    """The part of ``CaptionPredictor`` the adapter relies on."""

    decode_strategy: str
    beam_width: int
    length_penalty: float
    repetition_penalty: float
    no_repeat_ngram_size: int

    @property
    def config(self) -> _PredictorConfig: ...

    def predict_path(self, image_path: str | Path) -> str: ...


class CNNCaptioner(Captioner):
    """Wraps a loaded ``CaptionPredictor`` without changing how it decodes.

    Each image goes through ``predictor.predict_path``, one at a time, exactly
    as ``scripts/evaluate.py`` does. The recorded decode settings follow
    ``run_meta.json``: beam-only parameters are ``None`` for greedy decoding.
    """

    def __init__(self, predictor: PathPredictor, model: ComparedModelConfig) -> None:
        beam = predictor.decode_strategy == "beam"
        super().__init__(
            CaptionerIdentity(
                model_id=model.model_id,
                hub_repo=model.hub_repo,
                revision=model.revision,
                decode_settings={
                    "decode_strategy": predictor.decode_strategy,
                    "beam_width": predictor.beam_width if beam else None,
                    "length_penalty": predictor.length_penalty if beam else None,
                    "repetition_penalty": predictor.repetition_penalty,
                    "no_repeat_ngram_size": predictor.no_repeat_ngram_size,
                    "max_length": predictor.config.model.max_length,
                },
            )
        )
        self._predictor = predictor

    @classmethod
    def from_artifacts(
        cls,
        weights_path: str | Path,
        tokenizer_dir: str | Path,
        config: AppConfig,
        *,
        decode_strategy: DecodeStrategy | None = None,
        beam_width: int | None = None,
        length_penalty: float | None = None,
        repetition_penalty: float | None = None,
        no_repeat_ngram_size: int | None = None,
    ) -> CNNCaptioner:
        """Load the checkpoint with ``CaptionPredictor.from_artifacts``.

        Decoding arguments are passed through unchanged, so unset ones fall
        back to ``config.serve`` just as they do for ``scripts/evaluate.py``.
        The identity comes from ``config.compare.cnn``. TensorFlow is imported
        here, not when this module is imported.
        """
        from captioning.inference.predictor import CaptionPredictor

        predictor: Any = CaptionPredictor.from_artifacts(
            weights_path,
            tokenizer_dir,
            config,
            decode_strategy=decode_strategy,
            beam_width=beam_width,
            length_penalty=length_penalty,
            repetition_penalty=repetition_penalty,
            no_repeat_ngram_size=no_repeat_ngram_size,
        )
        return cls(predictor, config.compare.cnn)

    def _raw_captions(self, image_paths: list[Path]) -> list[str]:
        return [self._predictor.predict_path(path) for path in image_paths]