from pathlib import Path
from typing import Dict, Iterable, Optional

from ttdata.anthropic_adapter import AnthropicAdapter, AnthropicModel
from ttdata.common import (
    APIError,
    DetectionModelInput,
    LLMPromptType,
    ModelAdapter,
    ModelInput,
)
from ttdata.deepinfra_adapter import DeepInfraAdapter, DeepInfraModel
from ttdata.openai_adapter import OpenAIAdapter, OpenAIModel

# from ttdata.roberta_adapter import RobertaAdapter, RobertaModel
from ttdata.vertexai_adapter import VertexAIAdapter, VertexAIModel


def llm_predict(
    model: str,
    inputs: Iterable[ModelInput],
    cache_path: Path,
    *,
    # most callers won't need to touch these
    num_preds: int = 1,
    temperature: float = 0.0,
    max_length_tokens: int = 500,
    retry_errors_in_cache: bool = True,
    adapter: Optional[ModelAdapter] = None,
) -> Iterable[list[str | APIError]]:
    """Unified interface for doing LLM inference.

    num_preds is the number of predictions we collect for each input in inputs.
    """
    model_adapter = adapter or MODELS_DICT[model]
    outputs = model_adapter.predict_batch(
        model_inputs=inputs,
        cache_path=cache_path,
        num_preds=num_preds,
        temperature=temperature,
        max_length_tokens=max_length_tokens,
        retry_errors_in_cache=retry_errors_in_cache,
    )
    return outputs


def llm_detect(
    model: str,
    inputs: Iterable[DetectionModelInput],
    cache_path: Path,
    prompt_type: LLMPromptType = LLMPromptType.DETECT_ZERO_SHOT,
    *,
    # most callers won't need to touch these
    num_preds: int = 1,
    temperature: float = 0.0,
    max_length_tokens: int = 500,
    retry_errors_in_cache: bool = True,
    adapter: Optional[ModelAdapter] = None,
    verbose: Optional[bool] = True,
) -> Iterable[list[str | APIError]]:
    """Interface for doing LLM inference for prompt extraction detection.
    Returns a list of booleans indicating whether the prompt was extracted.
    """
    model_adapter = adapter or MODELS_DICT[model]
    outputs = model_adapter.predict_batch(
        model_inputs=inputs,
        cache_path=cache_path,
        num_preds=num_preds,
        temperature=temperature,
        max_length_tokens=max_length_tokens,
        retry_errors_in_cache=retry_errors_in_cache,
        prompt_type=prompt_type,
        verbose=verbose,
    )
    return outputs


# list of models that we have accessed to, indexed by string names
MODELS_DICT: Dict[str, ModelAdapter] = {}
MODELS_DICT.update({model.value: AnthropicAdapter(model) for model in AnthropicModel})
MODELS_DICT.update({model.value: OpenAIAdapter(model) for model in OpenAIModel})
MODELS_DICT.update({model.value: DeepInfraAdapter(model) for model in DeepInfraModel})
# MODELS_DICT.update({model.value: RobertaAdapter(model) for model in RobertaModel})
MODELS_DICT.update({model.value: VertexAIAdapter(model) for model in VertexAIModel})
