import spaces
import gradio as gr
from fpdf import FPDF
from maincode import *
import os
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    AutoModelForSequenceClassification
)
from huggingface_hub import InferenceClient
import ast
from datetime import datetime
import time

chat_history = []
nltk.download('punkt_tab')
token = os.environ.get('HF_TOKEN')
dataset_mng = DatasetManager(os.environ.get('DATASET'))
tokeniz = AutoTokenizer.from_pretrained("iarfmoose/t5-base-question-generator")
model = AutoModelForSeq2SeqLM.from_pretrained("iarfmoose/t5-base-question-generator")
qabertToken = AutoTokenizer.from_pretrained("iarfmoose/bert-base-cased-qa-evaluator")
qaBertModel = AutoModelForSequenceClassification.from_pretrained("iarfmoose/bert-base-cased-qa-evaluator")

class PDF(FPDF):
    def __init__(self, text):
        super().__init__()
        self.add_page()
        self.set_font("Arial", size=12)
        self.multi_cell(0, 10, text)

    def save(self, filename):
        self.output(filename)

api_enabled = True

def process_image(image):
    """
    Convert image to text
    """
    extracter = ExtractTextFromImage(image)
    extracter.process_file()
    return extracter.get_text()

def search_books(text, author, publisher, year, language, count):
    library = Books()
    library.search_book(text, author, publisher, year, language)
    books_result = library.get_result()
    html = ""
    for bok in books_result[:count]:
        html += generate_html(bok, library.get_pdf_link(bok))
    return html

def generate_qna_html(qna_list):
    html_code = """
    <style>
        .qna-container {
            margin: 10px 0;
            border: 1px solid #ccc;
            border-radius: 5px;
            padding: 10px;
        }
        .question {
            font-weight: bold;
            color: #333;
            margin-bottom: 5px;
            display: flex;
            align-items: center;
        }
        .arrow {
            margin-right: 10px;
        }
        .answer {
            margin-left: 20px;
            font-style: italic;
            color: #555;
        }
    </style>
    <div>
    """
    old = []
    for idx, qa in enumerate(qna_list):
        if qa['question'].lower() in old:
            continue
        old.append(qa['question'])
        question = qa['question']
        answer = qa['answer']
        html_code += f"""
        <div class="qna-container">
            <div class="question">
                <span class="arrow">▶</span>
                Q{idx + 1}. {question}
            </div>
            <div class="answer">{answer}</div>
        </div>
        """
    html_code += "</div>"
    return html_code

@spaces.GPU()
def get_qna(input_passage):
    """
    Generate question from input text
    """
    calcu = calculate_questions()
    calcu.process_paragraph(paragraph=input_passage)
    count = calcu.calculate()
    count = int(count / 2)
    qg = QAGenerator(
        tokeniz,
        model,
        tokenizerEV=qabertToken,
        modelEV=qaBertModel
    )
    qlist = qg.generate(input_passage, num_questions=count, answer_style="sentences")
    return generate_qna_html(qlist), qlist

def chat_with_huggingface(api_key, chat_message, history):
    """Function to send a chat message to the Hugging Face API and get a response using InferenceClient."""
    client = InferenceClient(api_key=api_key)
    messages = [{"role": "user", "content": chat_message}]
    for message in history:
        if "role" in message and "content" in message:
            messages.append(message)
    try:
        completion = client.chat.completions.create(
            model="Qwen/QwQ-32B-Preview", 
            messages=messages, 
            max_tokens=500
        )
        return completion.choices[0].message
    except Exception as e:
        return {"error": str(e)}

def search_question(search_query):
    result = dataset_mng.search(search_query, "train", "question")
    return result

def evaluate(passage, reference):
    evl = Evaluvate()
    return evl.get_result(passage, reference)

def chat_function(user_input, history):
    """Handle individual user chat, with integration to Hugging Face."""
    if history is None:
        history = []
    history.append({"role": "user", "content": user_input})
    api_response = chat_with_huggingface(token, user_input, history)
    if isinstance(api_response, dict) and "content" in api_response:
        bot_response_content = api_response["content"]
    else:
        bot_response_content = "Sorry, I couldn't understand that."
    bot_response = {"role": "assistant", "content": bot_response_content}
    history.append(bot_response)
    return history, ""

def generate_html(book_data, link):
    html_output = """
    <div style="font-family: Arial, sans-serif; line-height: 1.5;">
        <h2>{Title}</h2>
        <p><strong>Author:</strong> {Author}</p>
        <p><strong>Year:</strong> {Year}</p>
        <p><strong>Language:</strong> {Language}</p>
        <p><strong>Publisher:</strong> {Publisher}</p>
        <p><strong>Pages:</strong> {Pages}</p>
        <p><strong>Size:</strong> {Size}</p>
        <p><strong>Download Links:</strong></p>
        <ul>
            <li><a href="{link}" target="_blank">Download S1</a></li>
            <li><a href="{link2}" target="_blank">Download S2</a></li>
            <li><a href="{link3}" target="_blank">Download S3</a></li>
        </ul>
    </div>
    <hr style="border: 1px solid #ccc;">
    """.format(
        Title=book_data.get('Title', 'N/A'),
        Author=book_data.get('Author', 'N/A'),
        Year=book_data.get('Year', 'N/A'),
        Language=book_data.get('Language', 'N/A'),
        Publisher=book_data.get('Publisher', 'N/A'),
        Pages=book_data.get('Pages', 'N/A'),
        Size=book_data.get('Size', 'N/A'),
        link=link[0],
        link2=link[1],
        link3=link[2]
    )
    return html_output

def remove_question(question, qlist):
    if not isinstance(qlist, list):
        qlist = ast.literal_eval(qlist)
    for qa in qlist:
        if question.lower().strip() in qa.get("question").lower().strip():
            qlist.remove(qa)
    return generate_qna_html(qlist), qlist

def save_data_to_database(qlist, input_passage):
    if not isinstance(qlist, list):
        qlist = ast.literal_eval(qlist)
    to_save = []
    already_added = []
    for qa in qlist:
        try:
            result = dataset_mng.search(qa.get("question"), "train", "question", top_n=10)
            gr.Info(f"Question already exists in db : {qa.get('question')}")
        except:
            if qa.get("question").lower() in already_added:
                continue
            already_added.append(qa.get("question").lower())
            current_date_time = datetime.now()
            to_save.append(
                {
                        "question" : qa.get("question"),
                        "answer" : qa.get("answer"),
                        "refrence" : input_passage,
                        "data" : str(current_date_time.strftime("%d:%m:%Y %H:%M:%S")),
                        "timestamp" : int(time.time())
                }
            )
            continue
        for qas in result:
            if qas.get("match").lower() == qa.get("question").lower():
                continue
            else:
                if qa.get("question").lower() in already_added:
                    continue
                already_added.append(qa.get("question").lower())
                current_date_time = datetime.now()
                to_save.append(
                    {
                        "question" : qa.get("question"),
                        "answer" : qa.get("answer"),
                        "refrence" : input_passage,
                        "data" : str(current_date_time.strftime("%d:%m:%Y %H:%M:%S")),
                        "timestamp" : int(time.time())
                    }
                )
    if len(to_save) < 1:
        gr.Info("All question already exists in database")
    else:
        dataset_mng.add_to_dataset(to_save, token)
        gr.Info("Data saved to dataset")

def add_new_question(new_question, new_answer, qlist):
    if qlist == None:
        qlist = []
    if not isinstance(qlist, list):
        qlist = ast.literal_eval(qlist)
    qlist.append(
        {
            "question" : new_question,
            "answer" : new_answer
        }
    )
    return generate_qna_html(qlist), qlist

def edit_question_and_answer(real_question, Newquestion, Newanswer, qlist):
    if not isinstance(qlist, list):
        qlist = ast.literal_eval(qlist)
    for qa in qlist:
        if real_question.lower().strip() in qa.get("question").lower().strip():
            newqa = qa
            if Newquestion:
                newqa["question"] = Newquestion
            elif Newanswer:
                newqa["answer"] = Newanswer
            else:
                gr.Info("Nothing to update")
            qlist[qlist.index(qa)] = newqa
            break
    return generate_qna_html(qlist), qlist

with gr.Blocks() as demo:

    with gr.Tabs():

        with gr.Tab("Handwritten to Text"):
            gr.Markdown("OCR")
            image_input = gr.Image(label="Upload an Image", type="filepath")
            text_output = gr.Textbox(label="Extracted Text")
            process_button = gr.Button("Process Image")
            process_button.click(process_image, inputs=image_input, outputs=text_output)

        with gr.Tab("Search Books"):
            gr.Markdown("# Search For Books")
            text_input = gr.Textbox(label="Enter the book name (required)")
            author_input = gr.Textbox(label="Author (optional)", placeholder="Enter the author name")
            publisher_input = gr.Textbox(label="Publisher (optional)", placeholder="Enter the publisher name")
            year_input = gr.Textbox(label="Year (optional)", placeholder="Enter the publication year")
            language_input = gr.Textbox(label="Language (optional)", placeholder="Enter the language")
            result_count_slider = gr.Slider(label="Number of Results", minimum=1, maximum=25, step=1, value=5)
            pdf_links_output = gr.HTML(value="Boooks Result", label="Download Links")
            generate_pdf_button = gr.Button("Find Book")
            generate_pdf_button.click(
                search_books, 
                inputs=[text_input, author_input, publisher_input, year_input, language_input, result_count_slider], 
                outputs=pdf_links_output
            )

        with gr.Tab("Gen Q&A"):
            gr.Markdown("# Generate Question and Answer according to the paragraph")
            text_input_qna = gr.Textbox(label="Enter Paragraph")
            text_output_qna = gr.HTML(value="Q&A", label="Generated Q&A")
            generate_qna_button = gr.Button("Generate")
            text_output_qna_json = gr.Textbox(value="QNA_json", label="Generated Q&A Json", visible=False)
            generate_qna_button.click(get_qna, inputs=text_input_qna, outputs=[text_output_qna, text_output_qna_json])
            save_to_database = gr.Button("Save QNA")
            remove_question_input = gr.Textbox(label="Question To Remove")
            remove_q_button = gr.Button("Remove QNA")
            remove_q_button.click(remove_question, inputs=[remove_question_input, text_output_qna_json], outputs=[text_output_qna, text_output_qna_json])
            save_to_database.click(save_data_to_database, inputs=[text_output_qna_json, text_input_qna])
            gr.Markdown("""
                        ### To Edit the question you must have to enter the real question and here are thing how you can edit question and answer
                        - EDIT QUESTION : To edit question enter the real question first in the question to edit section, Then enter the question in the new question and then click on Edit Q&A button
                        - EDIT ANSWER : To edit answer do same thing enter the real question in the Question to edit textbox then add the new Answer in the New Answer text input and leave New Question empty and click on Edit Q&A button
                        - EDIT BOTH : To edit both thing ans and question just add real question in the question to edit and add new question in New Question and New answer in the New Answer and click on Edit Q&A to edit it.

                        ## These are the methods to edit the question and answer for the dataset
                        """)
            edit_question_real = gr.Textbox(label="Question To Edit", placeholder="Question to edit")
            with gr.Row():
                edit_question_input = gr.Textbox(label="New Question", placeholder="New Answer (Optional)")
                edit_answer_input = gr.Textbox(label="New Answer", placeholder="New Answer (Optional)")
            edit_qna_button = gr.Button("Edit Q&A")
            edit_qna_button.click(edit_question_and_answer, inputs=[edit_question_real, edit_question_input, edit_answer_input, text_output_qna_json], outputs=[text_output_qna, text_output_qna_json])
            new_question_input = gr.Textbox(label="New Question", placeholder="New question")
            new_answer = gr.Textbox(label="New Answer", placeholder="New Answer")
            button = gr.Button("Add Q&A")
            button.click(add_new_question, inputs=[new_question_input, new_answer], outputs=[text_output_qna, text_output_qna_json])

        with gr.Tab("Evaluate"):
            gr.Markdown("# Evaluate Texts")
            text_input_1 = gr.Textbox(label="Passage")
            text_input_2 = gr.Textbox(label="Reference")
            marks_input = gr.Number(label="Marks", interactive=False)
            hardness_range = gr.Slider(label="Hardness", minimum=1, maximum=10, step=1, value=5, interactive=False)
            html_output = gr.HTML(value="Evaluation text", label="Evaluation Result")
            evaluate_button = gr.Button("Evaluate")
            evaluate_button.click(evaluate, inputs=[text_input_1, text_input_2], outputs=html_output)
            def html_to_pdf(html_content):
                pdf = PDF(html_content)
                pdf.save("evaluation_output.pdf")
                return "evaluation_output.pdf"
            generate_pdf_button = gr.Button("Get PDF")
            pdf_output = gr.File(label="Download Evaluation PDF")
            generate_pdf_button.click(html_to_pdf, inputs=html_output, outputs=pdf_output)

        with gr.Tab("Search Question"):
            gr.Markdown("# Search Question")
            text_input_qna = gr.Textbox(label="Question")
            text_output_qna = gr.Textbox(value="Question And Answer result", label="Output")
            generate_qna_button = gr.Button("Find Answer")
            generate_qna_button.click(search_question, inputs=text_input_qna, outputs=text_output_qna)

        with gr.Tab("Chat"):
            gr.Markdown("# Chat Section")
            chat_display = gr.Chatbot(label="Chat", type="messages")
            with gr.Row():
                with gr.Column(scale=1):
                    chat_box = gr.Textbox(
                        label="Type your message here...",
                        placeholder="Start typing...",
                        interactive=True,
                    )
            chat_box.submit(chat_function, inputs=[chat_box, chat_display], outputs=[chat_display, chat_box])

demo.launch(show_api=False)