import PyPDF2
from openpyxl import load_workbook
from pptx import Presentation
import gradio as gr
import io
from huggingface_hub import InferenceClient
import re
import zipfile
import xml.etree.ElementTree as ET

# Constants
CHUNK_SIZE = 32000
SYSTEM_PROMPT = """
You are a helpful and informative assistant that can answer questions based on the content of documents. 
You will receive the content of a document and a question about it. 
Your task is to provide a concise and accurate answer to the question based solely on the provided document content.
If the document does not contain enough information to answer the question, simply state that you cannot answer the question based on the provided information.
"""

# Initialize the Mistral chat model
client = InferenceClient("mistralai/Mistral-Nemo-Instruct-2407")

def xml2text(xml):
    """Extracts text from XML data."""
    text = u''
    root = ET.fromstring(xml)
    for child in root.iter():
        text += child.text + " " if child.text is not None else ''
    return text

def extract_text_from_docx(docx_data, strip_content):
    """Extracts text from a DOCX file."""
    text = u''
    zipf = zipfile.ZipFile(io.BytesIO(docx_data))
    filelist = zipf.namelist()

    for fname in filelist:
        if re.match('word/header[0-9]*.xml', fname):
            text += xml2text(zipf.read(fname))
        elif re.match('word/footer[0-9]*.xml', fname):
            text += xml2text(zipf.read(fname))
    
    text += xml2text(zipf.read('word/document.xml'))
    zipf.close()

    if strip_content:
        text = strip_text(text)
    
    return f"{text}\n\n**Document Length:** {len(text)} characters"


def strip_text(text):
    """Strips unnecessary characters from text."""
    content = text.replace('\n', ' ')
    content = content.replace('\r', ' ')
    content = content.replace('\t', ' ')
    content = content.replace('  ', '')
    return content.strip()

def read_document(file, strip_content):
    """Reads the content of a document based on its file type."""
    file_path = file.name
    file_extension = file_path.split('.')[-1].lower()

    with open(file_path, "rb") as f:
        file_content = f.read()

    if file_extension == 'pdf':
        try:
            pdf_reader = PyPDF2.PdfReader(io.BytesIO(file_content))
            content = ''
            for page in range(len(pdf_reader.pages)):
                content += pdf_reader.pages[page].extract_text()
            if strip_content:
                content = strip_text(content)
            return content
        except Exception as e:
            return f"Error reading PDF: {e}"

    elif file_extension == 'xlsx':
        try:
            wb = load_workbook(io.BytesIO(file_content))
            content = ''
            for sheet in wb.worksheets:
                for row in sheet.rows:
                    for cell in row:
                        if cell.value is not None: 
                            content += str(cell.value) + ' '
            if strip_content:
                content = strip_text(content)
            return content
        except Exception as e:
            return f"Error reading XLSX: {e}"

    elif file_extension == 'pptx':
        try:
            presentation = Presentation(io.BytesIO(file_content))
            content = ''
            for slide in presentation.slides:
                for shape in slide.shapes:
                    if hasattr(shape, "text"):
                        content += shape.text + ' '
            if strip_content:
                content = strip_text(content)
            return content
        except Exception as e:
            return f"Error reading PPTX: {e}"

    elif file_extension == 'doc' or file_extension == 'docx':
        try:
            return extract_text_from_docx(file_content, strip_content)
        except Exception as e:
            return f"Error reading DOC/DOCX: {e}"

    else:
        try:
            content = file_content.decode('utf-8') 
            if strip_content:
                content = strip_text(content)
            return content
        except Exception as e:
            return f"Error reading file: {e}"


def split_content(content):
    """Splits content into chunks for processing."""
    chunks = []
    for i in range(0, len(content), CHUNK_SIZE):
        chunks.append(content[i:i + CHUNK_SIZE])
    return chunks


def chat_document(file, question, strip_content):
    """Handles chat with a document using Mistral."""
    content = str(read_document(file, strip_content))

    if len(content) > CHUNK_SIZE:
        content = content[:CHUNK_SIZE]

    message = f"""[INST] [SYSTEM] {SYSTEM_PROMPT} 
    Document Content: {content}
    Question: {question}
    Answer:"""

    stream = client.text_generation(message, max_new_tokens=4096, stream=True, details=True, return_full_text=False)
    output = ""
    for response in stream:
        if not response.token.text == "</s>":
            output += response.token.text
        yield output


def chat_document_v2(file, question, strip_content):
    """Handles chat with a document using Mistral and chunk-based approach."""
    content = str(read_document(file, strip_content))
    chunks = split_content(content)
    
    all_answers = []
    for chunk in chunks:
        message = f"""[INST] [SYSTEM] {SYSTEM_PROMPT} 
        Document Content: {chunk[:CHUNK_SIZE]}
        Question: {question}
        Answer:"""

        stream = client.text_generation(message, max_new_tokens=4096, stream=True, details=True, return_full_text=False)
        output = ""
        for response in stream:
            if not response.token.text == "</s>":
                output += response.token.text
        all_answers.append(output)

    # Summarize all answers using Mistral
    summary_prompt = """
    You are a helpful and informative assistant that can summarize multiple answers related to the same question. 
    You will receive a list of answers to a question, and your task is to generate a concise and comprehensive summary that incorporates the key information from all the answers.
    Avoid repeating information unnecessarily and focus on providing the most relevant and accurate summary based on the provided answers.
    
    Answers:
    """
    
    all_answers_str = "\n".join(all_answers)
    print(all_answers_str)
    summary_message = f"""[INST] [SYSTEM] {summary_prompt}
    {all_answers_str[:30000]}
    Summary:"""
    
    stream = client.text_generation(summary_message, max_new_tokens=4096, stream=True, details=True, return_full_text=False)
    output = ""
    for response in stream:
        if not response.token.text == "</s>":
            output += response.token.text
        yield output 

with gr.Blocks() as demo:
    with gr.Tabs():
        with gr.TabItem("Document Reader"):
            iface1 = gr.Interface(
                fn=read_document,
                inputs=[gr.File(label="Upload a Document"), gr.Checkbox(label="Strip Content", value=True)],
                outputs=gr.Textbox(label="Document Content"),
                title="Document Reader",
                description="Upload a document (PDF, XLSX, PPTX, TXT, CSV, DOC, DOCX and Code or text file) to read its content."
            )
        with gr.TabItem("Document Chat"):
            iface2 = gr.Interface(
                fn=chat_document,
                inputs=[gr.File(label="Upload a Document"), gr.Textbox(label="Question"), gr.Checkbox(label="Strip Content", value=True)],
                outputs=gr.Markdown(label="Answer"), 
                title="Document Chat",
                description="Upload a document and ask questions about its content."
            )
        with gr.TabItem("Document Chat V2"):
            iface3 = gr.Interface(
                fn=chat_document_v2,
                inputs=[gr.File(label="Upload a Document"), gr.Textbox(label="Question"), gr.Checkbox(label="Strip Content", value=True)],
                outputs=gr.Markdown(label="Answer"), 
                title="Document Chat V2",
                description="Upload a document and ask questions about its content (using chunk-based approach)."
            )

demo.launch()