"""
preprocessing.py
----------------
Converts raw knowledge-base samples into HuggingFace Datasets ready for
fine-tuning the multilingual-e5-large-instruct model.

Expected input schema (list of dicts)
--------------------------------------
[
    {
        "chunk_text": "The passage text ...",
        "docId": "doc-001",
        "questions": [
            {
                "augmented_questions": ["question a", "question b"],
                "noise_questions":     ["noise q 1"]
            },
            # ... more question groups
        ]
    },
    # ... more samples
]
"""

from __future__ import annotations

import logging
import string
from typing import Any

import numpy as np
import pandas as pd
from datasets import Dataset
from sklearn.model_selection import train_test_split

logger = logging.getLogger(__name__)


class MultilingualE5LargeInstructPreProcessing:
    """
    Transforms raw knowledge-base samples into train/test HuggingFace Datasets
    and an evaluator-compatible metadata dict.

    Parameters
    ----------
    task_description:
        The task string injected into every query prompt.
        Override this to adapt the model to a different retrieval task.
    test_size_ratio:
        Fraction of the data used for the test split (default ``0.1``).
    n_test_samples_per_chunk:
        Number of questions per chunk kept for testing (default ``10``).

    Example
    -------
    >>> from fiesta import MultilingualE5LargeInstructPreProcessing
    >>> preprocessor = MultilingualE5LargeInstructPreProcessing()
    >>> train_ds, test_ds, eval_data = preprocessor.preprocess(raw_data)
    """

    _DEFAULT_TASK = (
        "Given a web search query, retrieve relevant passages that answer the query"
    )

    def __init__(
        self,
        task_description: str | None = None,
        test_size_ratio: float = 0.1,
        n_test_samples_per_chunk: int = 10,
    ) -> None:
        self.task_description = task_description or self._DEFAULT_TASK
        self.test_size_ratio = test_size_ratio
        self.n_test_samples_per_chunk = n_test_samples_per_chunk

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def preprocess(
        self, data: list[dict[str, Any]]
    ) -> tuple[Dataset, Dataset, dict[str, Any]]:
        """
        Full preprocessing pipeline.

        Parameters
        ----------
        data:
            List of knowledge-base samples (see module docstring for schema).

        Returns
        -------
        train_dataset, test_dataset, evaluator_data
            Two HuggingFace Datasets and a metadata dict for the IR evaluator.
        """
        if not data:
            raise ValueError("Input data list is empty.")

        df = self._data_to_df(data)
        df = self._preprocess_queries(df)

        train_df, test_df = self._split_by_question(df)

        evaluator_data = self.get_evaluator_data(train_df, test_df)

        logger.info(
            "Preprocessing complete - train: %d rows, test: %d rows",
            len(train_df),
            len(test_df),
        )

        return (
            Dataset.from_pandas(train_df, preserve_index=False),
            Dataset.from_pandas(test_df, preserve_index=False),
            evaluator_data,
        )

    def get_evaluator_data(
        self,
        train: pd.DataFrame,
        test: pd.DataFrame,
    ) -> dict[str, Any]:
        """
        Build the three structures expected by
        :class:`sentence_transformers.evaluation.InformationRetrievalEvaluator`.

        Returns
        -------
        dict with keys:
            - ``train_test_chunks``      – ``{chunk_id: passage_text}``
            - ``test_queries``           – ``{question_id: query_text}``
            - ``testQueryId_to_chunkId`` – ``{question_id: [chunk_id]}``
        """
        train_test_chunks = self._get_all_chunks(train, test)
        test_queries: dict[int, str] = dict(zip(test["id"], test["anchor"]))
        testQueryId_to_chunkId: dict[int, list[int]] = dict(
            zip(
                test["id"].values,
                np.expand_dims(test["chunk_id"].values, 1).tolist(),
            )
        )
        return {
            "train_test_chunks": train_test_chunks,
            "test_queries": test_queries,
            "testQueryId_to_chunkId": testQueryId_to_chunkId,
        }

    def apply_prompt(self, query: str) -> str:
        """
        Wrap a raw query with the model's instruct-style prompt template.

        Parameters
        ----------
        query:
            Plain-text query string.

        Returns
        -------
        str
            Formatted prompt string.
        """
        return f"Instruct: {self.task_description}\nQuery: {query}"

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _data_to_df(self, data: list[dict[str, Any]]) -> pd.DataFrame:
        """Flatten raw samples into a tabular DataFrame."""
        if not data:
            raise ValueError("Input data list is empty.")
        records: list[dict[str, Any]] = []
        chunk_id = 0
        question_id = 0

        for sample in data:
            augmented = [
                q
                for question in sample["questions"]
                for q in question.get("augmented_questions", [])
            ]
            noise = [
                q
                for question in sample["questions"]
                for q in question.get("noise_questions", [])
            ]
            all_questions = augmented + noise

            chunk_text = sample["chunk_text"]
            doc_id = sample["docId"]
            n = len(all_questions)

            for i, q in enumerate(all_questions):
                records.append(
                    {
                        "id": question_id + i,
                        "anchor": q,
                        "positive": chunk_text,
                        "chunk_id": chunk_id,
                        "doc_id": doc_id,
                    }
                )

            question_id += n
            chunk_id += 1

        if not records:
            raise ValueError(
                "No question-chunk pairs could be extracted from the input data. "
                "Make sure each sample has at least one question with "
                "'augmented_questions' or 'noise_questions'."
            )

        return pd.DataFrame(records)

    def _preprocess_queries(self, data: pd.DataFrame) -> pd.DataFrame:
        """Lower-case, strip punctuation, then apply the instruct prompt."""
        _punct_table = str.maketrans("", "", string.punctuation)

        def _clean(text: str) -> str:
            return text.lower().translate(_punct_table)

        data = data.copy()
        data["anchor"] = data["anchor"].apply(_clean).apply(self.apply_prompt)
        return data

    def _split_by_question(
        self, data: pd.DataFrame
    ) -> tuple[pd.DataFrame, pd.DataFrame]:
        """
        For every unique chunk, hold out ``n_test_samples_per_chunk`` questions
        as the test set and keep the rest for training.

        Falls back to a simple random split if a chunk has too few questions.
        """
        unique_chunks = data["positive"].dropna().unique()

        if len(unique_chunks) < 2:
            raise ValueError(
                "Dataset must contain at least two distinct chunks to create "
                "a train/test split."
            )

        train_parts: list[pd.DataFrame] = []
        test_parts: list[pd.DataFrame] = []

        for chunk in unique_chunks:
            subset = data[data["positive"] == chunk]
            n = self.n_test_samples_per_chunk

            if len(subset) <= n:
                # Not enough samples - use the entire subset for training
                logger.warning(
                    "Chunk has only %d question(s) (≤ n_test_samples_per_chunk=%d). "
                    "Assigning all to train.",
                    len(subset),
                    n,
                )
                train_parts.append(subset)
            else:
                test_sample = subset.sample(n=n, random_state=42)
                train_sample = subset.drop(test_sample.index)
                train_parts.append(train_sample)
                test_parts.append(test_sample)

        train = pd.concat(train_parts, axis=0).reset_index(drop=True)
        test = (
            pd.concat(test_parts, axis=0).reset_index(drop=True)
            if test_parts
            else pd.DataFrame(columns=data.columns)
        )

        logger.info("Split - train: %s  test: %s", train.shape, test.shape)
        return train, test

    @staticmethod
    def _get_all_chunks(
        train: pd.DataFrame, test: pd.DataFrame
    ) -> dict[int, str]:
        """Merge train+test chunk pools, deduplicate, return as a dict."""
        combined = pd.concat(
            [train[["positive", "chunk_id"]], test[["positive", "chunk_id"]]]
        ).drop_duplicates(subset="chunk_id")
        return dict(zip(combined["chunk_id"], combined["positive"]))

    # ------------------------------------------------------------------
    # Legacy static helpers (kept for backwards compatibility)
    # ------------------------------------------------------------------

    @staticmethod
    def split_train_test_by_question(
        data: pd.DataFrame, test_size: float = 0.1
    ) -> tuple[pd.DataFrame, pd.DataFrame]:
        """Simple random train/test split by question row."""
        n_test = min(500, int(test_size * len(data)))
        return train_test_split(data, test_size=n_test)

    @staticmethod
    def split_train_test_by_chunk(
        data: pd.DataFrame,
        column: str = "positive",
        test_size: float = 0.1,
    ) -> tuple[pd.DataFrame, pd.DataFrame]:
        """Split unique chunk values into train/test, then filter rows."""
        unique_values = data[column].dropna().unique().tolist()
        if len(unique_values) < 2:
            raise ValueError(
                f"Column '{column}' has only one unique value; cannot split."
            )
        train_vals, test_vals = train_test_split(unique_values, test_size=test_size)
        return (
            data[data[column].isin(train_vals)],
            data[data[column].isin(test_vals)],
        )








