#!/user/bin/python
# -*- coding: utf-8 -*-
"""
Description: 文本切分算子
Create: 2023/11/09 10:17
"""

import random
from typing import List, Dict, Any

from loguru import logger

from deepdata.common.utils.text_splitter import TextSplitter
from deepdata.core.base_op import Slicer


class TextSegmentationOperator:
    def __init__(self, max_characters, chunk_size, chunk_overlap):
        try:
            self.text_splitter = TextSplitter(max_characters, chunk_size, chunk_overlap)
        except Exception as err:
            logger.exception(f"init text splitter failed, error is： {err}")
            raise Exception(83001, "init text splitter failed") from None

    def process(self, input_data: str) -> List[str]:
        if input_data.strip() == "":
            logger.info("input text is empty, return empty chunks.")
            return []
        return self.text_splitter.split_text(input_data)


class Segmentation(Slicer):
    """切片算法插件"""

    def __init__(self, *args, **kwargs):
        super(Segmentation, self).__init__(*args, **kwargs)
        self.max_characters = kwargs.get("maxCharacters", -1)
        self.chunk_size = kwargs.get("chunkSize", 800)
        self.chunk_overlap = kwargs.get("chunkOverlap", 100)
        self.slice_num = kwargs.get("sliceNum", -1)  # -1 means no limit, export all chunks
        self.text_splitter_op = TextSegmentationOperator(self.max_characters, self.chunk_size, self.chunk_overlap)
        self.last_ops = True

    def execute(self, sample: Dict[str, Any]) -> List[Dict]:

        try:
            chunks = self.text_splitter_op.process(sample[self.text_key])
        except Exception as err:
            logger.exception(f"split text failed, error is: {err}")
            raise Exception(83002, "split text failed") from None

        # Determine how many chunks to save
        num_to_save = len(chunks) if self.slice_num < 0 else min(self.slice_num, len(chunks))

        sample_list = []
        # If slice_num is set and less than total chunks, randomly sample; otherwise save all in order
        if 0 < self.slice_num < len(chunks):
            selected_indices = sorted(random.sample(range(len(chunks)), num_to_save))
        else:
            selected_indices = list(range(num_to_save))

        for idx in selected_indices:
            chunk = chunks[idx]
            temp_sample = {
                self.text_key: chunk,
                self.data_key: "",
                self.export_path_key: sample[self.export_path_key],
                self.filename_key: sample[self.filename_key],
                self.fileid_key: sample[self.fileid_key],
                "sequenceId": str(idx),
                "chunkSize": str(len(chunk)),
            }
            self.save_patch_sample(temp_sample, idx + 1, save_format="text")
            sample_list.append(temp_sample)

        sample["fileNum"] = len(chunks)
        # Keep original text if not all chunks are saved; clear text if all chunks are saved as slices
        if num_to_save < len(chunks):
            sample[self.text_key] = chunks[0] if chunks else ""
        else:
            # All chunks saved as slices, clear text to avoid duplication
            sample[self.text_key] = ""

        return sample_list