from typing import List, Union

from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from pydantic import BaseModel, Extra, validator

from mindsdb.integrations.handlers.rag_handler.settings import (
    DEFAULT_EMBEDDINGS_MODEL,
    RAGBaseParameters,
)

EVAL_COLUMN_NAMES = (
    "question",
    "answers",
    "context",
)

SUPPORTED_EVALUATION_TYPES = ("retrieval", "e2e")

GENERATION_METRICS = ("rouge", "meteor", "cosine_similarity", "accuracy")
RETRIEVAL_METRICS = ("cosine_similarity", "accuracy")


# todo make a separate class for evaluation parameters


class WriterLLMParameters(BaseModel):
    """Model parameters for the Writer LLM API interface"""

    writer_api_key: str
    writer_org_id: str = None
    base_url: str = None
    model_id: str = "palmyra-x"
    callbacks: List[StreamingStdOutCallbackHandler] = [StreamingStdOutCallbackHandler()]
    max_tokens: int = 1024
    temperature: float = 0.0
    top_p: float = 1
    stop: List[str] = []
    best_of: int = 5
    verbose: bool = False

    class Config:
        extra = Extra.forbid
        arbitrary_types_allowed = True


class WriterHandlerParameters(RAGBaseParameters):
    """Model parameters for create model"""

    llm_params: WriterLLMParameters
    generation_evaluation_metrics: List[str] = list(GENERATION_METRICS)
    retrieval_evaluation_metrics: List[str] = list(RETRIEVAL_METRICS)
    evaluation_type: str = "e2e"
    n_rows_evaluation: int = None  # if None, evaluate on all rows
    retriever_match_threshold: float = 0.7
    generator_match_threshold: float = 0.8
    evaluate_dataset: Union[List[dict], str] = "squad_v2_val_100_sample"

    class Config:
        extra = Extra.forbid
        arbitrary_types_allowed = True
        use_enum_values = True

    @validator("generation_evaluation_metrics", allow_reuse=True)
    def generation_evaluation_metrics_must_be_supported(cls, v):
        for metric in v:
            if metric not in GENERATION_METRICS:
                raise ValueError(
                    f"generation_evaluation_metrics must be one of {', '.join(str(v) for v in GENERATION_METRICS)}, got {metric}"
                )
        return v

    @validator("retrieval_evaluation_metrics", allow_reuse=True)
    def retrieval_evaluation_metrics_must_be_supported(cls, v):
        for metric in v:
            if metric not in GENERATION_METRICS:
                raise ValueError(
                    f"retrieval_evaluation_metrics must be one of {', '.join(str(v) for v in RETRIEVAL_METRICS)}, got {metric}"
                )
        return v

    @validator("evaluation_type", allow_reuse=True)
    def evaluation_type_must_be_supported(cls, v):
        if v not in SUPPORTED_EVALUATION_TYPES:
            raise ValueError(
                f"evaluation_type must be one of `retrieval` or `e2e`, got {v}"
            )
        return v

langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler
[((1785, 1845), 'pydantic.validator', 'validator', (['"""generation_evaluation_metrics"""'], {'allow_reuse': '(True)'}), "('generation_evaluation_metrics', allow_reuse=True)\n", (1794, 1845), False, 'from pydantic import BaseModel, Extra, validator\n'), ((2190, 2249), 'pydantic.validator', 'validator', (['"""retrieval_evaluation_metrics"""'], {'allow_reuse': '(True)'}), "('retrieval_evaluation_metrics', allow_reuse=True)\n", (2199, 2249), False, 'from pydantic import BaseModel, Extra, validator\n'), ((2591, 2637), 'pydantic.validator', 'validator', (['"""evaluation_type"""'], {'allow_reuse': '(True)'}), "('evaluation_type', allow_reuse=True)\n", (2600, 2637), False, 'from pydantic import BaseModel, Extra, validator\n'), ((857, 889), 'langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler', 'StreamingStdOutCallbackHandler', ([], {}), '()\n', (887, 889), False, 'from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n')]