"""helper utilities for ingesting documents"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/01_ingest.helpers.ipynb.

# %% auto 0
__all__ = ['CAPTION_DELIMITER', 'METADATA', 'FILE_METADATA', 'includes_caption', 'extract_tables', 'extract_files',
           'extract_extension', 'extract_file_dates', 'iso2date', 'date2iso', 'md5sum', 'get_mimetype',
           'extract_mimetype', 'is_random_plaintext', 'clean_text', 'doc_from_dict', 'dict_from_doc', 'create_document',
           'set_metadata_defaults', 'extract_file_metadata', 'ParagraphTextSplitter']

# %% ../../nbs/01_ingest.helpers.ipynb 3
from typing import List, Union, Optional
import warnings
import os
import os.path
import glob
from hashlib import md5
from functools import partial
import datetime
from langchain_core.documents import Document
from ..utils import contains_sentence

CAPTION_DELIMITER = '||CAPTION||'
def includes_caption(d:Document):
    """
    Returns True if content of supplied Document includes a table caption
    """
    table_captions = d.metadata.get('table_captions', '')
    if not table_captions: return False
    table_captions = table_captions.split(CAPTION_DELIMITER)
    for c in table_captions:
        if contains_sentence(c, d.page_content):
            return True
    return False


def extract_tables(filepath:Optional[str]=None, docs:Optional[List[Document]]=[]) -> List[Document]:
    """
    Extract tables from PDF and append to end of supplied Document list.
    Accepts either a `filepath` or a list of LangChain `Document` objects 
    all from a **single** file. If `filepath` is empty, the file path of interest 
    is extracted from `docs`. 

    Returns an updated list of Document objects appended with extracted tables.
    """
    # Lazy import PDFTables to avoid slow module loading on Windows
    from onprem.ingest.pdftables import PDFTables

    # get filepath of document under consideration
    docs = [] if not docs else docs
    if not filepath and not docs:
        raise ValueError('filepath and docs cannot both be empty.')
    if filepath and docs:
        raise ValueError('filepath and docs are mutually exclusive.')
    if docs:
        filepath = None if not docs else docs[0].metadata['source']
    if not filepath: return docs
        
    if extract_extension(filepath) != "pdf": return docs
    pdftab = PDFTables.from_file(filepath, verbose=False)
    md_tables = pdftab.get_markdown_tables()

    # tag document objects that contain extracted tables
    captions = pdftab.get_captions()
    for c in captions:
        c = c.strip()
        if len(c) < 8: continue
        for d in docs:
            if contains_sentence(c, d.page_content):
                table_captions = d.metadata.get('table_captions', [])
                if isinstance(table_captions, str):
                    table_captions = table_captions.split(CAPTION_DELIMITER)
                table_captions.append(c)
                d.metadata['table_captions'] = CAPTION_DELIMITER.join(table_captions)

    # augment docs with extracted tables
    tabledocs = []
    for md_table in md_tables:
        tabledoc = Document(page_content=md_table,
                metadata={'source':filepath, 'markdown':True, 'table':True})
        tabledocs.append(tabledoc)
    docs.extend(tabledocs)
    return docs


def extract_files(source_dir:str, follow_links=False, extensions:Optional[Union[dict,list]]=None):
    extensions = list(extensions.keys()) if isinstance(extensions, dict) else extensions
    if os.listdir(source_dir) == []:
        raise ValueError("%s: path is empty" % source_dir)
    walk = os.walk
    for root, _, filenames in walk(source_dir, followlinks=follow_links):
        for filename in filenames:
            if extensions and extract_extension(filename) not in extensions:
                continue
            try:
                yield os.path.join(root, filename)
            except:
                continue

def extract_extension(filepath:str, include_dot=False):
    """
    Extracts file extension (including dot) from file path
    """
    dot = "." if include_dot else ""
    return dot + filepath.rsplit(".", 1)[-1].lower()



def extract_file_dates(filepath):
    """
    Takes a file path and returns an ISO datetime string of last-modified
    and create date of file.

    Returns tuple of the form (create-date, last-modify-date)
    """
    from os.path import getmtime, getctime
    mtimestamp = getmtime(filepath)
    ctimestamp = getctime(filepath)
    mdate = datetime.datetime.fromtimestamp(mtimestamp).isoformat()
    cdate = datetime.datetime.fromtimestamp(ctimestamp).isoformat()
    return cdate, mdate

def iso2date(s):
    return datetime.datetime.fromisoformat(s)


def date2iso(d):
    return d.isoformat()


def md5sum(filepath):
    """
    Perform an MD5 hash of a file.
    """
    with open(filepath, "rb") as fobj:
        data = md5()
        for buf in iter(partial(fobj.read, 128), b""):
            data.update(buf)
    return data.hexdigest()

# %% ../../nbs/01_ingest.helpers.ipynb 4
import mimetypes
try:
    import magic

    get_magic = magic._get_magic_type(True)
    MAGIC_INSTALLED = True
except Exception as e:
    get_magic = None
    MAGIC_INSTALLED = False
    #warnings.warn(
        #"using mimetypes, since magic could not be imported (perhaps try python-magic-bin if on Windows):  %s"
        #% (e)
    #)

def get_mimetype(filepath):
    with open(filepath, "rb") as fobj:
        buf = fobj.read(4096)
    return get_magic.from_buffer(buf) if isinstance(get_magic, magic.Magic) else ""


def extract_mimetype(filepath):
    """
    Extract mimetype.
    Returns a tuple with extracted mimetype, type, subtype.
    """

    if MAGIC_INSTALLED:
        mtype = get_mimetype(filepath)
    else:
        mtype = mimetypes.guess_type(filepath)[0]
    mtype1 = None
    mtype2 = None
    if isinstance(mtype, str):
        try:
            mtype1, mtype2 = mtype.split("/")
            # Reading from buffer does not work for old office formats
            if "CDFV2" in mtype2 and isinstance(get_magic, magic.Magic):
                mtype = get_magic.from_file(filepath)
                if isinstance(mtype, str):
                    mtype1, mtype2 = mtype.split("/")
        except Exception:
            mtype = None
            mtype1 = None
            mtype2 = None
    return (mtype, mtype1, mtype2)


def is_random_plaintext(extension, mimetype):
    """
    Check mimetype for plain text
    """
    mimetype1, mimetype2 = mimetype.split('/')
    return mimetype1 == "text" or mimetype2 == "xml" or extension in {"txt"}

# %% ../../nbs/01_ingest.helpers.ipynb 5
def clean_text(text_s_or_b):
    """
    convert to string and strip.
    """
    text = (
        text_s_or_b.decode(errors="ignore")
        if isinstance(text_s_or_b, bytes)
        else text_s_or_b
    )
    text = text.strip()
    return text

# %% ../../nbs/01_ingest.helpers.ipynb 6
METADATA = {'source':None,
            'ocr':False,
            'table' : False,
            'markdown' :False,
            'page' :-1,
            'table_captions':'',
            'document_title' : '',
           }

def doc_from_dict(d:dict, content_field='page_content'):
    """
    Create LangChain Document from dicationary
    """
    metadata = d.copy()
    page_content = metadata.pop(content_field, '')
    return create_document(page_content, only_required_metadata=False, **metadata)


def dict_from_doc(doc, content_field = 'page_content'):
    """
    Create dictinoary from LangChain Document
    """
    metadata_with_content = doc.metadata.copy()
    metadata_with_content[content_field] = doc.page_content
    return metadata_with_content


def create_document(page_content:str, 
                    only_required_metadata:bool=True,
                   **kwargs):
    """
    Create document with required metadata keys from `METADATA`.
    """
    metadata = {k: v for k, v in kwargs.items() if k in METADATA} if only_required_metadata else kwargs
    if 'source' not in metadata: metadata['source'] = ''
    return Document(page_content=page_content, 
                    metadata=metadata)

def set_metadata_defaults(docs:List[Document], extra_keys:list=[]):
    """
    Sets Document metadata defaults
    """
    for doc in docs:
        # Keep only keys that are in METADATA or extra_keys
        doc.metadata = {k: v for k, v in doc.metadata.items() if k in METADATA or k in extra_keys}
        # Add keys from METADATA if they don't exist in doc.metadata
        doc.metadata.update({k: v for k, v in METADATA.items() if k not in doc.metadata})
    return docs


def _apply_file_callables(file_path:str, file_callables:dict):
    """
    Invokes file_callables on file path.

    Returns a dictionary with values containing results from callables for each key
    """
    if not file_callables: return {}

    if not os.path.exists(file_path):
        raise ValueError('file_path does not exist: {file_path}')

    results = {}
    for k,v in file_callables.items():
        results[k] = v(file_path)
    return results

FILE_METADATA = ['md5', 'mimetype', 'createdate', 'modifydate', 'extension']

def extract_file_metadata(file_path:str,
                          store_md5:bool=True,
                          store_mimetype:bool=True,
                          store_file_dates:bool=True,
                          file_callables:dict={}):
    """
    Extract file metadata
    """
    # extract metadata
    file_metadata = {}
    if store_md5:
        file_metadata['md5'] = md5sum(file_path)
    if store_mimetype:
        file_metadata['mimetype'], _, _ = extract_mimetype(file_path)
    if store_file_dates:
        file_metadata['createdate'], file_metadata['modifydate'] = extract_file_dates(file_path)
    ext = extract_extension(file_path)
    file_metadata['extension'] = ext
    file_metadata.update(_apply_file_callables(file_path, file_callables))
    return file_metadata



from typing import List
from langchain.text_splitter import TextSplitter
import re
class ParagraphTextSplitter(TextSplitter):
    def __init__(self, chunk_size: int = 5000, chunk_overlap: int = 0):
        super().__init__(chunk_size=chunk_size, chunk_overlap=chunk_overlap)

    def split_text(self, text: str) -> List[str]:
        paragraphs = re.split(r'\n\s*\n', text) # Paragraphs assumed to be separated by double newlines

        chunks = []
        for para in paragraphs:
            para = para.strip()
            if not para:
                continue
            if len(para) <= self._chunk_size:   # <- notice: _chunk_size (with underscore)
                chunks.append(para)
            else:
                chunks.extend(self._split_large_paragraph(para))
        return chunks

    def _split_large_paragraph(self, paragraph: str) -> List[str]:
        # Split a long paragraph into smaller chunks
        splits = []
        start = 0
        while start < len(paragraph):
            end = start + self._chunk_size
            splits.append(paragraph[start:end])
            start = end
        return splits

