import re

try:
    from nltk import sent_tokenize

    NLTK = True
except ImportError:
    NLTK = False

try:
    import chonkie

    CHONKIE = True
except ImportError:
    CHONKIE = False

from ..base import Pipeline

class Segmentation(Pipeline):
    def __init__(self, sentences=False, lines=False, Paragraphs=False, minlength=None, join=False, sections=False, cleantext=True, chunker=None, **kwargs):
        if not NLTK and sentences:
            raise ImportError('NLTK is not available - install "pipeline" extra to enable')
        
        if not CHONKIE and chunker:
            raise ImportError('Chonkie is not available - install "pipeline" extra to enable')

        self.sentences = sentences
        self.lines = lines
        self.paragraphs = Paragraphs
        self.sections = sections
        self.minlength = minlength
        self.join = join
        self.cleantext = cleantext

        self.chunker = self.createchunker(chunker, **kwargs) if chunker else None

    def __call__(self, text):
        texts = [text] if not isinstance(text, list) else text

        results = []
        for value in texts:
            value = self.text(value)
            results.append(self.parse(value))

        return results[0] if isinstance(text, str) else results
    
    def text(self, text):
        return text
    
    def parse(self, text):
        if self.chunker:
            content = [self.clean(x.text) for x in self.chunker(text)]
        elif self.sentences:
            content = [self.clean(x) for x in sent_tokenize(text)]
        elif self.lines:
            content = [self.clean(x) for x in re.split(r"\n{1,}", text)]
        elif self.paragraphs:
            content = [self.clean(x) for x in re.split(r"\n{2,}", text)]
        elif self.sections:
            split = r"\f" if "\f" in text else r"\n{3,}"
            content = [self.clean(x) for x in re.split(split, text)]
        else:
            content = self.clean(text)

        if isinstance(content, list):
            content = [x for x in content if x]
            return " ".join(content) if self.join else content
        
        return content
    
    def clean(self, text):
        if not self.cleantext:
            return text
        
        text = re.sub(r" +", " ", text)
        text = text.strip()

        return text if not self.minlength or len(text) >= self.minlength else None
    
    def createchunker(self, chunker, **kwargs):
        chunker = f"{chunker[0].upper() + chunker[1:]}Chunker"
        return getattr(chonkie, chunker)(**kwargs)