# @Author：青松

import lightning as L
from dsx.transformer.tokenizer import TransformerTokenizer
from torch.utils.data import Dataset, DataLoader


class TextTransform:
    def __init__(self, tokenizer: TransformerTokenizer, max_length=20):
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __call__(self, text):
        return self.tokenizer(
            text, max_length=self.max_length, padding=True, return_tensors=True
        )


class TextGenerationDataset(Dataset):
    def __init__(
        self,
        questions,
        answers,
        encoder_transform: TextTransform,
        decoder_transform: TextTransform,
    ):
        self.questions = questions
        self.answers = answers
        self.encoder_transform = encoder_transform
        self.decoder_transform = decoder_transform

    def __len__(self):
        return len(self.questions)

    def __getitem__(self, idx):
        question = self.questions[idx]
        answer = self.answers[idx]

        # 统一处理：始终为答案添加开始和结束标记
        processed_answer = (
            self.decoder_transform.tokenizer.bos_token
            + answer
            + self.decoder_transform.tokenizer.eos_token
        )

        question_encoded = self.encoder_transform(question)
        encoder_input_ids = question_encoded["input_ids"]
        encoder_pad_mask = question_encoded["attention_mask"]

        answer_encoded = self.decoder_transform(processed_answer)
        answer_ids = answer_encoded["input_ids"]
        answer_pad_mask = answer_encoded["attention_mask"]

        # 统一处理解码器输入输出序列
        decoder_input_ids = answer_ids[:-1].clone().detach()
        decoder_pad_mask = answer_pad_mask[:-1].clone().detach()
        decoder_target_ids = answer_ids[1:].clone().detach()

        return {
            "question": question,
            "answer": answer,
            "encoder_input_ids": encoder_input_ids,
            "encoder_pad_mask": encoder_pad_mask,
            "decoder_input_ids": decoder_input_ids,
            "decoder_pad_mask": decoder_pad_mask,
            "decoder_target_ids": decoder_target_ids,
        }

    @classmethod
    def from_file(
        cls,
        file_path,
        encoder_transform: TextTransform,
        decoder_transform: TextTransform,
    ):
        """
        从txt文件加载数据集
        txt格式应包含标签和文本，使用制表符分隔
        """

        questions = []
        answers = []

        # 读取txt文件
        with open(file_path, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip() == "":  # 跳过空行
                    continue
                try:
                    # 查找等号的位置，将行分割为问题和答案
                    idx = line.find("=")
                    # 将样本添加到列表中: (question, answer)，answer中去掉"="以及最后的"\n"
                    question = line[:idx]
                    answer = line[idx + 1 :].strip()
                    questions.append(question)
                    answers.append(answer)
                except Exception as e:
                    # 如果处理某行时出错，打印错误信息并跳过
                    print(f"Error processing line: {line}")
                    print(f"Error message: {e}")
                    continue

        # 创建数据集实例
        return cls(questions, answers, encoder_transform, decoder_transform)


class TextDataModule(L.LightningDataModule):
    def __init__(
        self,
        batch_size,
        encoder_transform: TextTransform,
        decoder_transform: TextTransform,
        train_data_file,
        val_data_file="",
        test_data_file="",
    ):
        super().__init__()

        self.batch_size = batch_size
        self.encoder_transform = encoder_transform
        self.decoder_transform = decoder_transform

        self.train_data_file = train_data_file
        self.val_data_file = val_data_file
        self.test_data_file = test_data_file

        self.test_dataset = None
        self.val_dataset = None
        self.train_dataset = None

    def prepare_data(self):
        # 下载或准备数据集的操作（如果需要）
        pass

    def setup(self, stage=None):
        # 加载训练数据集
        self.train_dataset = TextGenerationDataset.from_file(
            self.train_data_file,
            encoder_transform=self.encoder_transform,
            decoder_transform=self.decoder_transform,
        )

        # 加载验证数据集
        if self.val_data_file == "":
            self.val_dataset = self.train_dataset
        else:
            self.val_dataset = TextGenerationDataset.from_file(
                self.val_data_file,
                encoder_transform=self.encoder_transform,
                decoder_transform=self.decoder_transform,
            )

        # 加载测试数据集
        if self.test_data_file == "":
            self.test_dataset = self.train_dataset
        else:
            self.test_dataset = TextGenerationDataset.from_file(
                self.test_data_file,
                encoder_transform=self.encoder_transform,
                decoder_transform=self.decoder_transform,
            )

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=self.batch_size)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size)
