import gradio as gr
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch
import transformers
from transformers import pipeline
import re
# import HTML
from pypdf import PdfReader

def remove_references(text):
    text = re.sub(r'\[\d+\]', '', text) ##[ref]
    text = re.sub(r'\[https?://[^\[\]]+\s[^\[\]]+\]', '', text) ##hyperlink with text
    text = re.sub(r'\[https?://[^\[\]]+\]', '', text) ##just the hyperlink
    # text = html.unescape(text)
    text = re.sub(r'\s+', ' ', text).strip() ##clear out the white spaces
    return text

    
def extract_text_from_pdf(file_path):
  reader = PdfReader(file_path)
  text = ""
  for page in reader.pages[2:]:
      text += page.extract_text() + "\n"

  return text  

      
def model(model_name):
  tokenizer = AutoTokenizer.from_pretrained(model_name)
  model = AutoModelForQuestionAnswering.from_pretrained(model_name,return_dict = False)
  model_pipeline = pipeline(
    "question-answering",
    model = model,
    tokenizer = tokenizer
    )

  return model_pipeline

def qa_result(context, question, file_path):
    model_name = "timpal0l/mdeberta-v3-base-squad2"
    pipe = model(model_name)
    if file_path is not None:
        context = extract_text_from_pdf(file_path)
        result = pipe(question=question, context=context)
        answered = result['answer']
        text = remove_references(answered)
        
    elif len(context) == 0 and len(question) == 0:
            text = "Որպեսզի ես կարողանամ քեզ օգնել, ինձ պիտի տրամադրես համապատասխան տեքստն ու հարցերը։"
    elif len(context) == 0:
        text = "Ես չեմ կարողանամ քեզ օգնել եթե ինձ չտրամադրես տեքստը"
    elif len(question) == 0:
        text = "Ես չեմ կարողանամ քեզ օգնել եթե ինձ չտաս հարցդ"
    else:
        # if file_path is not None:
        #     # File was uploaded, extract text from the file
        #     context = extract_text_from_pdf(file_path)
        # else: No file uploaded, use the provided context as is
        
        result = pipe(question=question, context=context)
        answered = result['answer']
        text = remove_references(answered)
        text = text.replace('(', '', 1)
        text = text.replace(',', '', len(text)-1)
    return text.capitalize()

theme = gr.themes.Soft().set(
    body_background_fill='*background_fill_secondary',
    body_text_color_subdued='*body_text_color',
    body_text_color_subdued_dark='*chatbot_code_background_color'
)
def add_file(history, file):
    history = history + [((file.name,), None)]
    return history


app = gr.Interface(
    fn=qa_result,
    btn=gr.UploadButton("📁", file_types=[".pdf", ".csv", ".doc"], ),
    inputs=['textbox', 'text', 'file'],
    outputs='textbox',
    title='Ողջու՛յն։ Ես քո արհեստական բանականությամբ օգնականն եմ',
    theme=theme,
    description='Տու՛ր ինձ տեքստ, ու տեքստին վերաբերող հարցեր, ու ես կօգնեմ քեզ պատասխանել հարցերին'
)


app.launch(inline=False)