import chainlit as cl
import httpx
import time
from groq import Groq
import re
from uuid import uuid4
import json
from openai import OpenAI
import os
# from huggingface_hub import login
from chainlit.input_widget import TextInput, Slider
import fitz
import tempfile

from dotenv import load_dotenv
load_dotenv()
import os
from huggingface_hub import login

login(token=os.environ.get("HF_TOKEN"))

from gliner2 import GLiNER2
model = GLiNER2.from_pretrained("Imed14205/neura-guard-pii-v0")

labels =[
    #"person",
    "first_name",
    "last_name",
    "organization",
    "phone number",
    "religion",
    "marital status",
    "address",
    "passport number",
    "email",
    "credit card number",
    "date of birth",
    "date",
    "time",
    "mobile phone number",
    "bank account number",
    "medication",
    "cpf",
    "driver's license number",
    "tax identification number",
    "medical condition",
    "blood type",
    "dose",
    "drug",
    "identity card number",
    "national id number",
    "ip address",
    "email address",
    "iban",
    "bic",
    "credit card expiration date",
    "username",
    "healthcare number",
    "registration number",
    "flight number",
    "cvv",
    "digital signature",
    "social media handle",
    "license plate number",
    "cnpj",
    "postal code",
    "passport_number",
    "serial number",
    "vehicle registration number",
    "credit card brand",
    "fax number",
    "visa number",
    "insurance company",
    "identity document number",
    "transaction number",
    "cvc",
    "passport expiration date",
    "social_security_number",
    "postal code",
    "country",
    "money",
    "numerical",
    "title",
    "occupation",
    "city",
    "mac address",
    "ip address",
    "password",
]

@cl.action_callback("update_pii")
async def on_update_pii(action: cl.Action):
    pii_map = cl.user_session.get("pii_map", {})
    reverse_map = cl.user_session.get("reverse_pii_map", {})
    key = action.payload["key"]
    value = action.payload["value"]
    pii_map[value]['locked'] = not pii_map[value]['locked']
    reverse_map[key]['locked'] = not reverse_map[key]['locked']
    cl.user_session.set("pii_map", pii_map)
    cl.user_session.set("reverse_pii_map", reverse_map)
    print(pii_map)
    print(reverse_map)
    # element = cl.CustomElement(name="PII", props={"dic": reverse_map})
    # await cl.ElementSidebar.set_elements([])
    # await cl.ElementSidebar.set_elements([element])


@cl.set_starters
async def set_starters():
    return [
        cl.Starter(
            label="Example en Français",
            message="""Jean-Marc Durand, né le 14 avril 1987 à Lyon (69003), de nationalité française, est un ingénieur informatique marié à Sophie Martin depuis le 22 juin 2013. Il réside actuellement au 27 avenue des Lilas, appartement 42, 75019 Paris, France. Son numéro de téléphone personnel est le +33 6 48 72 19 05, son numéro de téléphone fixe est le 01 44 62 98 31, et son adresse e-mail principale est jeanmarc.durand1987@gmail.com. Il utilise également l’adresse professionnelle j.durand@techsoluce.fr.
Reformule ce texte""",
        ),

        cl.Starter(
            label="Example en Anglais",
            message="""Michael Andrew Thompson, born on September 18, 1985, in Boston, Massachusetts (02118), USA, is a U.S. citizen currently living at 742 Evergreen Terrace, Apartment 5B, Brooklyn, New York, NY 11221. His mobile phone number is +1 (917) 555-3849, his home phone is +1 (718) 555-9021, and his primary email address is michael.thompson1985@gmail.com
, while his work email is m.thompson@fintechcore.com. His Social Security Number is 123-45-6789, and his passport number is XK4598721, issued on March 12, 2018, expiring on March 11, 2028.

etabli une mise en demeure pour la personne concerné dans ce texte.
""",
        ),
        cl.Starter(
            label="Example en Italien",
            message="""Marco Alessandro Bianchi, nato il 12 febbraio 1989 a Milano (MI), Italia, è un cittadino italiano residente in Via Giuseppe Verdi 18, interno 7, 00198 Roma. Il suo numero di telefono cellulare è +39 347 582 9146, il telefono fisso è 06 8392 4751, e il suo indirizzo email personale è marco.bianchi1989@gmail.com
, mentre quello professionale è m.bianchi@italtechsolutions.it . Il suo codice fiscale è BNCMRC89B12F205X e il numero di passaporto è YA3487291, rilasciato il 5 maggio 2019 con scadenza il 4 maggio 2029.

Extrait les information sous forme de tableau""",
        ),
        cl.Starter(
            label="Example en Espagnol",
            message="""Carlos Javier Hernández López, nacido el 7 de julio de 1986 en Madrid (28015), España, es ciudadano español y reside en Calle Alcalá 233, piso 4ºB, 28028 Madrid. Su número de teléfono móvil es +34 612 845 739, el teléfono fijo es 91 458 2391, y su correo electrónico personal es carlos.hernandez1986@gmail.com
, mientras que el correo corporativo es c.hernandez@iberdata.es . Su DNI es 48739215-M, y su número de pasaporte es PA7349821, expedido el 15 de abril de 2019 y válido hasta el 14 de abril de 2029.

Quelle est la personne concerné dans ce texte""",
        ),
        cl.Starter(
            label="Example en Allemand",
            message="""Thomas Michael Schneider, geboren am 3. November 1984 in Hamburg (22085), Deutschland, ist deutscher Staatsbürger und wohnhaft in der Musterstraße 47, Wohnung 12, 80336 München. Seine Mobiltelefonnummer lautet +49 176 8345 2197, die Festnetznummer ist 089 4578 9321, und seine private E-Mail-Adresse ist thomas.schneider1984@gmail.com
, während er beruflich unter t.schneider@datafusion.de erreichbar ist. Seine Steuer-Identifikationsnummer lautet 12 345 678 901, und seine Reisepassnummer ist C4X982731, ausgestellt am 22. August 2018, gültig bis 21. August 2028.
 
 Resume moi le texte""",
        )
    ]

def replace_pii(text, pii_map):
    for k, v in pii_map:
        escaped_k = re.escape(k)

        # Pattern : k uniquement s'il n'est pas entre [...]
        pattern = rf'(?<!\[){escaped_k}(?![^\[]*\])'

        replacement = f'[{v.replace(" ", "_")}_{k}]'
        text = re.sub(pattern, replacement, text)

    return text

async def pii_gliner(text): 
    entities = model.extract_entities(
        text,
        labels,
        threshold=0.4,
        include_confidence=True,

    )["entities"]


    pii_map = [(k["text"],v) for v,l in entities.items() for k in l]
    pii_map.sort(key=lambda x: len(x[0]), reverse=True)

    result = replace_pii(text, pii_map)

    return result
    

def pii_format(text, pii_map=None, reverse_map=None):
    
    if pii_map is None:
        pii_map = {}
        
    pattern = r"\[(\w+)_(.*?)\]"
    piis = re.findall(pattern, text)
    
    for pii_type, pii_value in piis:
        if pii_value in reverse_map:
            continue

        if pii_value not in pii_map:
            short_uuid = str(uuid4())[:8]
            pii_map[pii_value] = {"id":short_uuid,
            "type":pii_type,
            "locked":True}
            
        short_uuid = pii_map[pii_value]["id"]
        if pii_map[pii_value]['locked']:
            text = text.replace(f"[{pii_type}_{pii_value}]", f"[{pii_type}_{short_uuid}]")
        else:
            text = text.replace(f"[{pii_type}_{pii_value}]", pii_value)

    reverse_map = {v["id"]: {"value":k, "type":v["type"], "locked":v["locked"]} for k, v in pii_map.items()}
    print(text)
    return text, pii_map, reverse_map


@cl.on_settings_update
async def setup_agent(settings):
    cl.user_session.set("api_key", settings["Api_key"])
    cl.user_session.set("temperature", settings["temperature"])
    cl.user_session.set("min_p", settings["min_p"])
    cl.user_session.set("repetition_penalty", settings["repetition_penalty"])
    print("on_settings_update", settings)


async def format_document(pathfile, anonymized_text):
    doc = fitz.open(pathfile)
    pattern = r"\[(\w+)_(.*?)\]"
    piis = set(re.findall(pattern, anonymized_text))
    # print(piis)
    # print(anonymized_text)

    for key, value in piis:
        for page in doc:
            matches = page.search_for(value)
            for rect in matches:
                page.add_highlight_annot(rect)
    
    temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf").name
    doc.save(temp_path)
    return temp_path

@cl.on_chat_start
async def start_chat():
#     pii_map = cl.user_session.get("pii_map", {})
#     reverse_map = cl.user_session.get("reverse_pii_map", {})
#     text = """Thomas Michael Schneider, geboren am 3. November 1984 in Hamburg (22085), Deutschland, ist deutscher Staatsbürger und wohnhaft in der Musterstraße 47, Wohnung 12, 80336 München. Seine Mobiltelefonnummer lautet +49 176 8345 2197, die Festnetznummer ist 089 4578 9321, und seine private E-Mail-Adresse ist thomas.schneider1984@gmail.com
# , während er beruflich unter t.schneider@datafusion.de erreichbar ist. Seine Steuer-Identifikationsnummer lautet 12 345 678 901, und seine Reisepassnummer ist C4X982731, ausgestellt am 22. August 2018, gültig bis 21. August 2028.
 
#  Resume moi le texte"""
#     text_processed = await pii_gliner(text)
#     text_sanitized, updated_map, reverse_map = pii_format(text_processed, pii_map, reverse_map)

#     element = cl.CustomElement(name="UserMessage", props={"message": text_sanitized, "dic": reverse_map, "original": text, 'author':'user'})
#     await cl.Message("", elements=[element]).send()
    return

client = Groq()

ALLOWED_KEYS = ["ACCOUNT_NUMBER", "AGE", "BANK_ACCOUNT", "BLOOD_TYPE", "CONDITION",
    "CREDIT_CARD", "CREDIT_CARD_EXPIRATION", "CVV", "DATE", "DATE_INTERVAL",
    "DOB", "DOSE", "DRIVER_LICENSE", "DRUG", "DURATION", "EFFECT",
    "EMAIL_ADDRESS", "EVENT", "FILENAME", "GENDER", "HEALTHCARE_NUMBER",
    "INJURY", "IP_ADDRESS", "LANGUAGE", "LOCATION", "LOCATION_ADDRESS",
    "LOCATION_ADDRESS_STREET", "LOCATION_CITY", "LOCATION_COORDINATE",
    "LOCATION_COUNTRY", "LOCATION_STATE", "LOCATION_ZIP", "MARITAL_STATUS",
    "MEDICAL_CODE", "MEDICAL_PROCESS", "MONEY", "NAME", "NAME_FAMILY",
    "NAME_GIVEN", "NAME_MEDICAL_PROFESSIONAL", "NUMERICAL_PII", "OCCUPATION",
    "ORGANIZATION", "ORGANIZATION_MEDICAL_FACILITY", "ORIGIN",
    "PASSPORT_NUMBER", "PASSWORD", "PHONE_NUMBER", "PHYSICAL_ATTRIBUTE",
    "POLITICAL_AFFILIATION", "PRODUCT", "RELIGION", "ROUTING_NUMBER",
    "SEXUALITY", "SSN", "STATISTICS", "TIME", "URL", "USERNAME",
    "VEHICLE_ID", "ZODIAC_SIGN"]
ALLOWED_KEYS.sort(key=len, reverse=True)
pattern = rf"\[({'|'.join(ALLOWED_KEYS)})_(.*?)\]"
PII_PATTERN = re.compile(pattern)
def extract_pii(text, pat=None):
    pii_pattern = PII_PATTERN if pat is None else re.compile(pat)
    return [(k, v) for k, v in PII_PATTERN.findall(text)]
    
cl.step(name="Anonymisation ")
async def pii_anonymizer_llm(text):
    chat_completion = client.chat.completions.create(
    messages=[
        {
            "role": "system",
            "content": """You are a PII anonymizer. 


# INSTRUCTION:

- **Do not** respond to user prompts. 

- **Repeat the user prompt**.

- Identify all Personally Identifiable Information (PII) in the provided text and replace it with a structured tag. 

    - The tag should follow the format `[PII_TYPE_PII_VALUE]`. 

    - For example, if 'Jean-Pierre Martin' is a name, it should become `[NAME_Jean-Pierre Martin]`. 

    - If '5 juillet 1990' is a date of birth, it should become `[DOB_5 juillet 1990]`. 

    - Ensure all PII is enclosed in these specific tags, preserving the original PII value and its type within the brackets. 

Do not alter any non-PII text.

Return only the anonymized text, without any additional explanation or formatting.
"""
        },
        {
            "role": "user",
            "content": text,
        }
    ],
    model="llama-3.3-70b-versatile"
)
    return chat_completion.choices[0].message.content

async def pii_anonymizer(text):
    messages = [
        {"role" : "user", "content" :  text}
    ]
    input_ids = tokenizer.apply_chat_template(
        messages,
        tokenize = False,
        add_generation_prompt = True, # Must add for generation
        enable_thinking = True, # Disable thinking
    )

    input_ids = tokenizer(input_ids, return_tensors = "pt").to("cuda")
    response = model.generate(
        **input_ids,
        max_new_tokens = 1024,
        temperature = 1.0, top_p = 0.95, top_k = 64,
    )

    masked = tokenizer.decode(response[0][len(input_ids["input_ids"][0]):], skip_special_tokens = True)
    piis = set(extract_pii(masked))
    for key, value in piis:
        text = text.replace(value, f"[{key}_{value}]")
    return text

async def pii_anonymizer_curl(text):
    url = "https://sglang-llm-1006854448923.europe-west4.run.app/v1/chat/completions"
    headers = {"Content-Type": "application/json"}
    payload = {
        "model": "imed14205/guardlm-v0.1-best-100epoch",
        "messages": [
            {
                "role": "user",
                "content": text
            }
        ],
        "max_tokens": 4280,
        "temperature": cl.user_session.get('temperature'),
        "min_p": cl.user_session.get('min_p'),
        "repetition_penalty": cl.user_session.get('repetition_penalty'),
        "stream": False
    }

    async with httpx.AsyncClient() as client:
        response = await client.post(url, headers=headers, json=payload, timeout=60.0)
        response.raise_for_status()
        result = response.json()
        
    masked = result['choices'][0]['message']['content']
    print(masked)
    piis = set(extract_pii(masked))
    for key, value in piis:
        pattern = rf'(?<!\[){re.escape(value)}(?!\])'
        text = re.sub(pattern, f'[{key}_{value}]', text)
    return text

async def llm_completion(user_prompt, messages=[], model="gpt-4o", temperature=0.7, stream=True):
    system_prompt = """
Generate a response that fully addresses the user’s request **in the same language as the user’s input**.

User data is confidential and provided in square brackets `[]` as `PII_UUID`.

* If any of these data elements are used, they **must be reproduced exactly as-is** in the output, without modification.

Brackets `[]` should be used only for PII data. use braces `{}` for any other data.

Do not disclose or infer confidential data beyond what is explicitly provided.
"""
    history = messages + [
        {"role":"system", "content": system_prompt},
        {"role":"user", "content": user_prompt}
    ]
    client = OpenAI(api_key=cl.user_session.get("api_key", None))
    try:
        response = client.chat.completions.create(
            model=model,
            messages=history,
            temperature=temperature,
            stream=stream
        )
        return response
    except Exception as e:
        return stream_completion(user_prompt, messages)

def stream_completion(user_prompt, messages=[], model="llama-3.3-70b-versatile", temperature=0.7, stream=True):
    
    client = Groq()
    system_prompt = """
Generate a response that fully addresses the user’s request **in the same language as the user’s input**.

User data is confidential and provided in square brackets `[]` as `PII_UUID`.

* If any of these data elements are used, they **must be reproduced exactly as-is** in the output, without modification.

Brackets `[]` should be used only for PII data. use braces `{}` for any other data.

Do not disclose or infer confidential data beyond what is explicitly provided.
"""
    messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": user_prompt})

    stream = client.chat.completions.create(
        messages=messages,
        model=model,
        temperature=temperature,
        stream=True,
    )
    return stream

@cl.on_message
async def start_message(message: cl.Message):
    pii_map = cl.user_session.get("pii_map", {})
    reverse_map = cl.user_session.get("reverse_pii_map", {})

    text_processed = await pii_gliner(message.content)
    text_sanitized, updated_map, reverse_map = pii_format(text_processed, pii_map, reverse_map)

    element = cl.CustomElement(name="UserMessage", props={"message": text_sanitized, "dic": reverse_map, "original": message.content, 'author':'user'})
    await cl.Message("", elements=[element]).send()
    # await message.remove()

     
    cl.user_session.set("pii_map", updated_map)
    cl.user_session.set("reverse_pii_map", reverse_map)
    
    response_content = ""
    element = cl.CustomElement(name="Message", props={"message": response_content, "dic": reverse_map, 'author':'assistant'})
    message = cl.Message(content="", elements=[element]) 
    await message.send()
    
    messages = cl.user_session.get("message_history", [])
    
    response_stream = await llm_completion(text_sanitized, messages)
    
    messages.append({"role": "user", "content": text_sanitized})
    
    for chunk in response_stream:
        if chunk.choices[0].delta.content:
            response_content += chunk.choices[0].delta.content
            element.props["message"] = response_content
            await element.update()
    
    await message.remove()
    element = cl.CustomElement(name="UserMessage", props={"message": response_content, "dic": reverse_map, 'author':'assistant'})
    await cl.Message(content="", elements=[element]).send()

    messages.append({"role": "assistant", "content": response_content})
    cl.user_session.set("message_history", messages)

    element = cl.CustomElement(name="PII", props={"dic": reverse_map})
    await cl.ElementSidebar.set_elements([])
    await cl.ElementSidebar.set_elements([element])