import itertools
import logging
import os
import tempfile

import pandas as pd

from autorag.evaluator import Evaluator
from autorag.utils import cast_qa_dataset, cast_corpus_dataset, validate_qa_from_corpus_dataset

logger = logging.getLogger("AutoRAG")


class Validator:
    def __init__(self, qa_data_path: str, corpus_data_path: str):
        """
        Initialize a Validator object.

        :param qa_data_path: The path to the QA dataset.
            Must be parquet file.
        :param corpus_data_path: The path to the corpus dataset.
            Must be parquet file.
        """
        # validate data paths
        if not os.path.exists(qa_data_path):
            raise ValueError(f"QA data path {qa_data_path} does not exist.")
        if not os.path.exists(corpus_data_path):
            raise ValueError(f"Corpus data path {corpus_data_path} does not exist.")
        if not qa_data_path.endswith('.parquet'):
            raise ValueError(f"QA data path {qa_data_path} is not a parquet file.")
        if not corpus_data_path.endswith('.parquet'):
            raise ValueError(f"Corpus data path {corpus_data_path} is not a parquet file.")
        self.qa_data = pd.read_parquet(qa_data_path, engine='pyarrow')
        self.corpus_data = pd.read_parquet(corpus_data_path, engine='pyarrow')
        self.qa_data = cast_qa_dataset(self.qa_data)
        self.corpus_data = cast_corpus_dataset(self.corpus_data)

    def validate(self, yaml_path: str, qa_cnt: int = 5, random_state: int = 42):
        # sample QA data
        sample_qa_df = self.qa_data.sample(qa_cnt, random_state=random_state)
        sample_qa_df.reset_index(drop=True, inplace=True)

        # get doc_id
        temp_qa_df = sample_qa_df.copy(deep=True)
        flatten_retrieval_gts = temp_qa_df['retrieval_gt'].apply(
            lambda x: list(itertools.chain.from_iterable(x))).tolist()
        target_doc_ids = list(itertools.chain.from_iterable(flatten_retrieval_gts))

        # make sample corpus data
        sample_corpus_df = self.corpus_data.loc[self.corpus_data['doc_id'].isin(target_doc_ids)]
        sample_corpus_df.reset_index(drop=True, inplace=True)

        validate_qa_from_corpus_dataset(sample_qa_df, sample_corpus_df)

        # start Evaluate at temp project directory
        with tempfile.NamedTemporaryFile(suffix='.parquet') as qa_path, \
                tempfile.NamedTemporaryFile(suffix='.parquet') as corpus_path, \
                tempfile.TemporaryDirectory() as temp_project_dir:
            sample_qa_df.to_parquet(qa_path.name, index=False)
            sample_corpus_df.to_parquet(corpus_path.name, index=False)

            evaluator = Evaluator(qa_data_path=qa_path.name, corpus_data_path=corpus_path.name,
                                  project_dir=temp_project_dir)
            evaluator.start_trial(yaml_path)

        logger.info("Validation complete.")
