import logging
import re
from typing import Literal, List, Tuple
from langchain_text_splitters import MarkdownTextSplitter, LatexTextSplitter, RecursiveCharacterTextSplitter

# ========== Interface ==========
class ResponseGetChunk:
    text: str
    source: str

# ========== Main Class ==========
class Structured():
    def __init__(self):
        pass

    @staticmethod
    def get_chunks(buffer: bytearray) -> Tuple[List[ResponseGetChunk], int]:
        text = clean_crlf(buffer.decode("utf-8"))
        char_count = len(text)

        rows = csv_to_array(text)   
        header = rows[0]
        content_index = header.index("content")

        if content_index == -1:
            logging.error('Content column not found')
            return [], -1
            
        chunks = [
            row[content_index]
            for row in rows[1:]
            if row[content_index]
        ]

        return [
            {
                "text": chunk,
                "source": "structured"
            }
            for chunk in chunks
        ], char_count

    @staticmethod
    def split_file(buffer: bytearray, file_extension: Literal["md", 'txt', 'tex', 'csv']) -> Tuple[List[ResponseGetChunk], int]:
        whole_file = clean_new_line(clean_crlf(buffer.decode("utf-8")))
        char_count = len(whole_file)

        if file_extension == 'md':
            return Structured.split_markdown(whole_file), char_count
        elif file_extension == 'txt' or file_extension == 'csv':
            return Structured.split_text(whole_file), char_count
        elif file_extension == 'tex':
            return Structured.split_latex(whole_file), char_count
        else:
            logging.error(f"Unsupported file extension: {file_extension}")
            return [], -1

    @staticmethod
    def split_markdown(text: str) -> List[ResponseGetChunk]:
        splitter = MarkdownTextSplitter(chunk_size=1200, chunk_overlap=200)
        split_result = splitter.split_text(text)
        output = [
            {
                "text": chunk,
                "source": "structured"
            }
            for chunk in split_result
        ]
        return output

    @staticmethod
    def split_text(text: str) -> List[ResponseGetChunk]:
        splitter = RecursiveCharacterTextSplitter(chunk_size=1200, chunk_overlap=200)
        split_result = splitter.split_text(text)
        output = [
            {
                "text": chunk,
                "source": "structured"
            }
            for chunk in split_result
        ]
        return output
    
    @staticmethod
    def split_latex(text: str) -> List[ResponseGetChunk]:
        splitter = LatexTextSplitter(chunk_size=1200, chunk_overlap=200)
        split_result = splitter.split_text(text)
        output = [
            {
                "text": chunk,
                "source": "structured"
            }
            for chunk in split_result
        ]
        return output

# ========== Function ==========
def clean_crlf(text: str) -> str:
    return re.sub(r'\r\n?|\n', '\n', text)

def clean_new_line(text: str) -> str:
    return re.sub(r'\n{3,}', '\n\n', text)

def csv_to_array(text: str) -> list[list[str]]:
    p = ""
    row = [""]
    ret = [row]
    i = 0
    r = 0
    s = True

    for l in text:
        if l == '"':
            if s and l == p:
                row[i] += l
            s = not s
        elif l == "," and s:
            l = ""
            row.append(l)
            i += 1
        elif l == "\n" and s:
            if p == "\r":
                row[i] = row[i][:-1]
            row = [""]
            ret.append(row)
            r += 1
            i = 0
        else:
            row[i] += l
        p = l