import gradio as gr

from transformers import T5Tokenizer, T5ForConditionalGeneration
from mailparser import parse_from_file
from bs4 import BeautifulSoup
from gliner import GLiNER
from typing import Dict, Union, List

import spacy
import re
import os
import en_core_web_sm
nlp = en_core_web_sm.load()

t5_tokenizer = T5Tokenizer.from_pretrained("t5-small")
t5_model = T5ForConditionalGeneration.from_pretrained("t5-small")

_MODEL = {}
_CACHE_DIR = os.environ.get("CACHE_DIR", None)

def accept_mail(file_path):
    email = parse_from_file(file_path)
    return email

def clean_email(email):
    soup = BeautifulSoup(email.body, 'html.parser')
    for tag in soup.find_all(['style', 'link']):
        tag.decompose()
    cleaned_text = ' '.join(soup.get_text(separator=' ').split())
    return cleaned_text

def remove_special_characters(text):
    pattern = r'[=_-]+'
    cleaned_text = re.sub(pattern, '', text)
    return cleaned_text

def get_sentences(further_cleaned_text):
    doc = nlp(further_cleaned_text)
    sentences = [sent.text for sent in doc.sents]
    return sentences

def get_model(model_name: str = None, multilingual: bool = False):
    if model_name is None:
        model_name = "urchade/gliner_base" if not multilingual else "urchade/gliner_multilingual"
    if model_name not in _MODEL:
        _MODEL[model_name] = GLiNER.from_pretrained(model_name, cache_dir=_CACHE_DIR)
    return _MODEL[model_name]

def parse_query(sentences, labels, threshold=0.3, nested_ner=False, model_name=None, multilingual=False):
    model = get_model(model_name, multilingual)
    results = []
    for sentence in sentences:
        _entities = model.predict_entities(sentence, labels, threshold=threshold)
        results.extend([{"text": entity["text"], "label": entity["label"]} for entity in _entities])
    return results

def refine_entities_with_t5(entities):
    inputs = "refine entities: " + " ; ".join([f"{entity['text']} as {entity['label']}" for entity in entities])
    input_ids = t5_tokenizer.encode(inputs, return_tensors="pt", add_special_tokens=True)
    outputs = t5_model.generate(input_ids)
    result = t5_tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result

def present(email_file, labels, multilingual=False):
    email = accept_mail(email_file)
    cleaned_text = clean_email(email)
    further_cleaned_text = remove_special_characters(cleaned_text)
    sentence_list = get_sentences(further_cleaned_text)
    entities = parse_query(sentence_list, labels, threshold=0.3, nested_ner=False, model_name="urchade/gliner_base", multilingual=multilingual)
    refined_entities = refine_entities_with_t5(entities)
    email_info = {
        "Subject": email.subject,
        "From": email.from_,
        "To": email.to,
        "Date": email.date,
        "Extracted Entities": entities,  # Prepare entities for DataFrame if needed
        "Refined Entities": refined_entities
    }
    return [email_info[key] for key in ["Subject", "From", "To", "Date", "Extracted Entities", "Refined Entities"]]

labels = ["PERSON", "PRODUCT", "DEAL", "ORDER", "ORDER PAYMENT METHOD", "STORE", "LEGAL ENTITY", "MERCHANT", "FINANCIAL TRANSACTION", "UNCATEGORIZED", "DATE"]

demo = gr.Interface(
    fn=present,
    inputs=[
        gr.components.File(label="Upload Email (.eml file)"),
        gr.components.CheckboxGroup(choices=labels, label="Labels to Detect", value=labels),
        gr.components.Checkbox(label="Use Multilingual Model")
    ],
    outputs=[
        gr.components.Textbox(label="Subject"),
        gr.components.Textbox(label="From"),
        gr.components.Textbox(label="To"),
        gr.components.Textbox(label="Date"),
        gr.components.Dataframe(headers=["Text", "Label"], label="Extracted Entities"),
        gr.components.Textbox(label="Refined Entities")
    ],
    title="Email Info Extractor",
    description="Upload an email file (.eml) to extract its details and detected entities."
)
demo.launch(share=True)