import re
import gradio as gr
from entity_dictionary import ENTITY_DICTIONARY


# tokenize
def tokenize(text):
    return re.findall(r"\w+|[.!?]", text)


# dictionary matching
def find_dictionary_match(tokens, start):
    for category, phrases in ENTITY_DICTIONARY.items():
        if category == "REGEX_PATTERNS":
            continue

        for phrase in phrases:
            p_tokens = phrase.split()
            span = tokens[start:start + len(p_tokens)]

            if len(span) == len(p_tokens) and \
               [t.lower() for t in span] == [p.lower() for p in p_tokens]:
                return category, len(p_tokens), " ".join(span)

    return None


# dictionary entity detection
def detect_dictionary_entities(text):
    tokens = tokenize(text)
    entities = []

    i = 0
    sentence_start = True

    while i < len(tokens):
        token = tokens[i]

        if token in ".!?":
            sentence_start = True
            i += 1
            continue

        if not token.isalpha():
            sentence_start = False
            i += 1
            continue

        # sentence start
        if sentence_start:
            match = find_dictionary_match(tokens, i)

            if match:
                cat, length, value = match
                entities.append({"text": value, "type": cat})
                i += length
            else:
                i += 1

            sentence_start = False
            continue

        # capitalized groups
        if token[0].isupper():
            match = find_dictionary_match(tokens, i)

            if match:
                cat, length, value = match
                entities.append({"text": value, "type": cat})
                i += length
                continue

            phrase = [token]
            j = i + 1

            while (
                j < len(tokens)
                and tokens[j].isalpha()
                and tokens[j][0].isupper()
            ):
                phrase.append(tokens[j])
                j += 1

            entities.append({
                "text": " ".join(phrase),
                "type": "CAP_GROUP"
            })

            i = j
            continue

        # normal dictionary matching
        match = find_dictionary_match(tokens, i)

        if match:
            cat, _, value = match
            entities.append({"text": value, "type": cat})

        i += 1

    return entities


# regex detection
def detect_regex_entities(text, existing_entities):
    if "REGEX_PATTERNS" not in ENTITY_DICTIONARY:
        return []

    occupied = set()

    for e in existing_entities:
        for m in re.finditer(re.escape(e["text"]), text, re.IGNORECASE):
            occupied.update(range(m.start(), m.end()))

    matches = []

    for category, pattern in ENTITY_DICTIONARY["REGEX_PATTERNS"].items():
        for m in re.finditer(pattern, text):

            span = set(range(m.start(), m.end()))

            # avoid overlap
            if span & occupied:
                continue

            matches.append({
                "text": m.group(),
                "type": category,
                "start": m.start(),
                "end": m.end(),
                "length": m.end() - m.start()
            })

    matches.sort(key=lambda x: x["length"], reverse=True)

    return [{"text": m["text"], "type": m["type"]} for m in matches]


# masking
def mask_text(text, entities):
    masked = text
    entity_map = {}
    counter = 1

    for e in sorted(entities, key=lambda x: len(x["text"]), reverse=True):

        ph = f"z{counter}"
        counter += 1

        masked = re.sub(
            rf"\b{re.escape(e['text'])}\b",
            ph,
            masked,
            flags=re.IGNORECASE
        )

        entity_map[ph] = e["text"]

    return masked, entity_map


# main function
def detect_and_mask(text):

    if not text.strip():
        return "", "", ""

    # option B entity detection
    dict_entities = detect_dictionary_entities(text)
    regex_entities = detect_regex_entities(text, dict_entities)

    entities = dict_entities + regex_entities

    # remove duplicates
    unique = {}
    for e in entities:
        key = e["text"].lower()

        if key not in unique:
            unique[key] = e

    entities = list(unique.values())

    # mask
    masked_text, entity_map = mask_text(text, entities)

    # format entities
    if entities:
        entity_output = "\n".join(
            [f"{e['text']} --> {e['type']}" for e in entities]
        )
    else:
        entity_output = "No entities detected."

    # format placeholders
    mapping_output = "\n".join(
        [f"{k} --> {v}" for k, v in entity_map.items()]
    )

    return entity_output, masked_text, mapping_output


# gradio interface
iface = gr.Interface(
    fn=detect_and_mask,
    inputs=gr.Textbox(
        lines=4,
        placeholder="Enter sentence..."
    ),
    outputs=[
        gr.Textbox(lines=10, label="Detected Entities"),
        gr.Textbox(lines=4, label="Masked Sentence"),
        gr.Textbox(lines=10, label="Placeholder Mapping")
    ],
    title="Entity Detection and Masking",
    description="Detects entities and generates masked sentence",
    api_name="detect_entities"
)

iface.launch()