import pandas as pd
import os
from utils.text_splitter import srt_chunker as chun
from utils.vector_db import VectorDB
import fitz
from datetime import datetime
import re

def load_document(file_path, raw_text=False):
    file_extension = os.path.splitext(file_path)[1].lower()
    if file_extension == '.pdf':
        return load_pdf(file_path)
    elif file_extension == '.csv':
        return load_csv(file_path)
    elif file_extension == '.srt':
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
        paragraphs = content.split('\n\n')
        documents = [{"text": paragraph, "paragraph_number": index + 1}
                     for index, paragraph in enumerate(paragraphs) if paragraph.strip()]
        return documents
    elif file_extension == '.md':
        if raw_text:
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            return content
        else:
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            paragraphs = content.split('\n\n')
            documents = [{"text": paragraph, "paragraph_number": index + 1}
                         for index, paragraph in enumerate(paragraphs) if paragraph.strip()]
            return documents
    else:
        raise ValueError(f"Unsupported file type: {file_extension}")

def load_pdf(file_path):
    """
    Load a PDF file and extract its text content into paragraphs.

    Args:
        file_path (str): The path to the PDF file.

    Returns:
        list[dict]: A list of dictionaries, each containing a paragraph of text and its paragraph number.
    """
    paragraphs = []
    pdf_document = fitz.open(file_path)

    for page in pdf_document:
        output = page.get_text("blocks")
        previous_block_id = None

        for block in output:
            if block[6] == 0:  # Check if the block is a text block
                cleaned_output = block[4].strip().replace("\n", " ")
                if cleaned_output and len(cleaned_output) >= 5:
                    if previous_block_id != block[5]:
                        paragraphs.append(cleaned_output)
                    else:
                        paragraphs[-1] += " " + cleaned_output
                    previous_block_id = block[5]

    # Adding the paragraph number to each paragraph
    documents = [{"text": paragraph, "paragraph_number": index + 1}
                 for index, paragraph in enumerate(paragraphs)]

    return documents


def load_csv(file_path):
    try:
        df = pd.read_csv(file_path)
        documents = []
        if "Paragraph text" in df.columns:
            df = df.rename(columns={"Paragraph text": "text"})
            df["page_number"] = df.index + 1
            documents = df.to_dict(orient="records")
            return documents
        else:
            for index, row in df.iterrows():
                content = row.to_json()
                documents.append({"text": content, "page_number": (index+1)})
            return documents
    except Exception as e:
        print(e)
        return []


def upload_document_to_vectorstore(chunks: list[dict], vector_db: VectorDB):
    """
    Uploads the chunks and their metadata to the vector database in batches.

    Args:
        chunks (list[dict]): A list of dictionaries, each containing a chunk of the document.
        vector_db (VectorDB): An instance of the VectorDB class to interact with the vector database.

    Returns:
        list: A list of points created in the vector store.
    """
    batch_size = 2000
    qdrant_batch_size = 500
    print(f"Uploading {len(chunks)} chunks to the vector store.")
    embeddings = []
    for batch_start in range(0, len(chunks), batch_size):
        batch_chunks = chunks[batch_start: batch_start + batch_size]
        batch_chunks = [chunk["text"] for chunk in batch_chunks]
        print(f"Encoding batch of {len(batch_chunks)} chunks.")
        batch_embeddings = vector_db.batch_encode_text(batch_chunks)

        embeddings.extend(batch_embeddings)

    points = vector_db.create_points(embeddings, chunks)
    print(f"Uploading {len(points)} points to the vector store.")

    for batch_start in range(0, len(points), qdrant_batch_size):
        batch_points = points[batch_start: batch_start + qdrant_batch_size]
        vector_db.insert_data(batch_points)
    print("Upload complete.")
    return points


def document_chunker(document: list[dict], chunking_function) -> list[dict]:
    """
    Chunk a document into smaller pieces using a provided chunking function.

    Args:
        document (list[dict]): A list of dictionaries, each containing a 'text' key with the document content.
        chunking_function (callable): A function that takes a string and returns a list of chunks.

    Returns:
        list[dict]: A list of dictionaries, each containing a chunk of the document.
    """
    chunks = []
    for page in document:
        text = page.pop("text")
        text_chunks = chunking_function(text)
        for chunk in text_chunks:
            chunks.append({"text": chunk, **page})

    return chunks


def transcript_chunker(document: list[dict], filename) -> list[dict]:
    """
    Chunk a transcript document into smaller pieces.

    Args:
        document (list[dict]): A list of dictionaries, each containing a 'text' key with the transcript content.

    Returns:
        list[dict]: A list of dictionaries, each containing a chunk of the transcript.
    """

    chunks = []

    def transcript_formatter(
        transcript): return f"{transcript['timestamp']} {transcript['speaker']}: {transcript['text']}"
    count = 0
    for page in document:
        text = page.pop("text")
        transcript = chun(text)

        if transcript is None:
            if "File name" in text:
                continue
            if chunks:
                chunks[-1]["text"] += text
            continue

        # Add the filename to the transcript
        transcript["filename"] = filename

        # Adding the index of the transcript to the transcript
        transcript["index"] = count
        count += 1

        # Add the current transcript to the chunks
        chunks.append(transcript)

        # Add the previous message to the current message
        if len(chunks) > 1:
            chunks[-1]["previous_message"] = transcript_formatter(chunks[-2])

        # Add the next message to the previous message
        if len(chunks) > 1:
            chunks[-2]["next_message"] = transcript_formatter(chunks[-1])

    return chunks

def merge_chunks(chunks: list[dict]) -> list[dict]:
    """
    Merge consecutive chunks that have the same speaker

    Args:
        chunks (list[dict]): A list of dictionaries, each containing a chunk of text and its metadata.

    Returns:
        list[dict]: A list of merged chunks.
    """
    if not chunks:
        return []

    merged_chunks = []
    current_chunk = chunks[0]
    for chunk in chunks[1:]:
        if chunk["speaker"] == current_chunk["speaker"]:
            # Merge the text of the current chunk with the next one
            current_chunk["text"] += " " + chunk["text"]
        else:
            # Add the current chunk to the merged list and start a new one
            merged_chunks.append(current_chunk)
            current_chunk = chunk
    # Add the last chunk
    merged_chunks.append(current_chunk)
    return merged_chunks


def generate_filename_from_details(minute_block_md: str) -> str:
    """
    Generates a filename from a minute block by extracting details 
    from a line like "| 24 March 2025, Guoqi Wu, AVP CSD | |".
    Format: YYYY-MM-DD_Name1Name2_RolePart1_RolePart2_CLE.md
    """
    lines = minute_block_md.splitlines()
    details_str = None

    if len(lines) > 1:
        # The line containing details for the filename is typically the second line of the block
        # e.g., | 24 March 2025, Guoqi Wu, AVP CSD | |
        header_info_line = lines[1] 
        
        # Extract content from between the first and second pipe characters
        line_parts = header_info_line.split('|')
        if len(line_parts) > 2 and line_parts[1].strip():
            details_str = line_parts[1].strip()
        else:
            raise ValueError(f"Could not parse details for filename from line: '{header_info_line}' in minute block.")
    else:
        raise ValueError("Minute block has insufficient lines to extract details for filename.")

    if not details_str:
        raise ValueError("Failed to extract details string from minute block.")

    # Regex to separate date from the rest of the details
    match = re.match(r"(\d{1,2}\s+[A-Za-z]+\s+\d{4}),\s*(.+)", details_str)
    if not match:
        raise ValueError(f"Details string '{details_str}' does not match expected format 'Date, Details...'")

    date_part_str = match.group(1) 
    name_and_role_parts_str = match.group(2)

    # Convert date to YYYY-MM-DD
    try:
        dt_object = datetime.strptime(date_part_str, "%d %B %Y")
        formatted_date = dt_object.strftime("%Y-%m-%d")
    except ValueError:
        try:
            dt_object = datetime.strptime(date_part_str, "%d %b %Y")
            formatted_date = dt_object.strftime("%Y-%m-%d")
        except ValueError as e:
            raise ValueError(f"Date string '{date_part_str}' could not be parsed: {e}")

    parts = [part.replace(" ", "") for part in name_and_role_parts_str.split(", ")]
    combined_name_role = "_".join(parts)
    combined_name_role_sanitized = re.sub(r'[^\w-]', '', combined_name_role.replace("(", "").replace(")", ""))

    return f"{formatted_date}_{combined_name_role_sanitized}_CLE.md", (formatted_date, combined_name_role_sanitized)

def load_prompt(prompt_file_name:str, working_directory) -> str:
    """
    Load a prompt from a file in the specified working directory's prompts folder.

    Args:
        prompt_file_name (str): The name of the prompt file to load.
        working_directory (str): The directory where the prompt file is located.

    Returns:
        str: The content of the prompt file.
    """
    prompt_file_path = os.path.join(working_directory, "prompts", prompt_file_name+".txt")
    if not os.path.exists(prompt_file_path):
        raise FileNotFoundError(f"Prompt file '{prompt_file_name}' at path '{prompt_file_path}' does not exist.")

    with open(prompt_file_path, 'r', encoding='utf-8') as file:
        return file.read()