import ast
import json
import re

import gradio as gr
import spaces
import torch

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer


# ==========================================================
# Model Configuration
# ==========================================================

BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
ADAPTER_ID = "mirajbhandari/Entity_Extcation_Quen"

SYSTEM_PROMPT = (
    "You are an NER model. Extract named entities from the sentence and "
    'return ONLY a JSON list of objects with keys "text" and "type". '
    "Allowed types: PERSON, ORGANIZATION, LOCATION, DATE, EVENT, PRODUCT, "
    "MONEY, TIME, WORK_OF_ART, LANGUAGE, NORP, FAC, GPE."
)


# ==========================================================
# Load Tokenizer
# ==========================================================

print("Loading tokenizer...")

try:
    tokenizer = AutoTokenizer.from_pretrained(ADAPTER_ID)
except Exception:
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id


# ==========================================================
# Global Model Variable
# ==========================================================

model = None


# ==========================================================
# Entity Colors
# ==========================================================

ENTITY_COLORS = {
    "PERSON": "#fecaca",
    "ORGANIZATION": "#bfdbfe",
    "LOCATION": "#bbf7d0",
    "GPE": "#a7f3d0",
    "DATE": "#fde68a",
    "TIME": "#fed7aa",
    "EVENT": "#ddd6fe",
    "PRODUCT": "#fbcfe8",
    "MONEY": "#c7d2fe",
    "WORK_OF_ART": "#e9d5ff",
    "LANGUAGE": "#bae6fd",
    "NORP": "#f5d0fe",
    "FAC": "#d9f99d",
}


# ==========================================================
# Load Model on ZeroGPU
# ==========================================================

def load_model():
    global model

    if model is not None:
        return model

    print("Loading base model on GPU...")

    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        dtype=torch.float16,
        device_map="cuda",
        low_cpu_mem_usage=True,
    )

    print("Loading LoRA adapter...")

    model = PeftModel.from_pretrained(
        base_model,
        ADAPTER_ID,
    )

    print("Merging LoRA adapter...")

    model = model.merge_and_unload()
    model.eval()

    print("Model is ready!")

    return model


# ==========================================================
# Prompt
# ==========================================================

def build_messages(sentence):
    return [
        {
            "role": "system",
            "content": SYSTEM_PROMPT,
        },
        {
            "role": "user",
            "content": sentence,
        },
    ]


# ==========================================================
# Parse Model Output
# ==========================================================

def parse_entities(model_output):
    output = model_output.strip()

    # Remove Markdown code block if returned by the model.
    output = re.sub(
        r"^```(?:json)?\s*",
        "",
        output,
        flags=re.IGNORECASE,
    )

    output = re.sub(
        r"\s*```$",
        "",
        output,
    )

    start = output.find("[")
    end = output.rfind("]")

    if start == -1 or end == -1 or end < start:
        raise ValueError("The model did not return a valid JSON list.")

    json_text = output[start:end + 1]

    try:
        entities = json.loads(json_text)
    except json.JSONDecodeError:
        entities = ast.literal_eval(json_text)

    if not isinstance(entities, list):
        raise ValueError("The model result must be a JSON list.")

    cleaned_entities = []

    for entity in entities:
        if not isinstance(entity, dict):
            continue

        text = str(entity.get("text", "")).strip()
        entity_type = str(entity.get("type", "")).strip().upper()

        if text and entity_type:
            cleaned_entities.append(
                {
                    "text": text,
                    "type": entity_type,
                }
            )

    return cleaned_entities


# ==========================================================
# Find Entity Positions
# ==========================================================

def find_entity_spans(sentence, entities):
    spans = []
    occupied_positions = []

    for entity in entities:
        entity_text = entity["text"]
        entity_type = entity["type"]

        # Exact match first.
        matches = list(
            re.finditer(
                re.escape(entity_text),
                sentence,
            )
        )

        # Case-insensitive match if exact matching fails.
        if not matches:
            matches = list(
                re.finditer(
                    re.escape(entity_text),
                    sentence,
                    flags=re.IGNORECASE,
                )
            )

        for match in matches:
            start = match.start()
            end = match.end()

            overlaps = any(
                start < existing_end and end > existing_start
                for existing_start, existing_end in occupied_positions
            )

            if overlaps:
                continue

            spans.append(
                {
                    "start": start,
                    "end": end,
                    "text": sentence[start:end],
                    "type": entity_type,
                }
            )

            occupied_positions.append((start, end))
            break

    spans.sort(key=lambda item: item["start"])

    return spans


# ==========================================================
# Create Highlighted Text
# ==========================================================

def create_highlighted_output(sentence, spans):
    if not sentence:
        return []

    if not spans:
        return [(sentence, None)]

    output = []
    current_position = 0

    for span in spans:
        start = span["start"]
        end = span["end"]

        if start > current_position:
            output.append(
                (
                    sentence[current_position:start],
                    None,
                )
            )

        output.append(
            (
                sentence[start:end],
                span["type"],
            )
        )

        current_position = end

    if current_position < len(sentence):
        output.append(
            (
                sentence[current_position:],
                None,
            )
        )

    return output


# ==========================================================
# Prediction
# ==========================================================

@spaces.GPU(duration=120)
@torch.inference_mode()
def predict(sentence):
    sentence = sentence.strip()

    if not sentence:
        raise gr.Error("Please enter a sentence.")

    try:
        current_model = load_model()

        messages = build_messages(sentence)

        prompt = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )

        inputs = tokenizer(
            prompt,
            return_tensors="pt",
        )

        inputs = {
            key: value.to("cuda")
            for key, value in inputs.items()
        }

        generated_ids = current_model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=False,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

        # Remove prompt tokens from generated output.
        generated_tokens = generated_ids[
            :,
            inputs["input_ids"].shape[1]:
        ]

        raw_output = tokenizer.batch_decode(
            generated_tokens,
            skip_special_tokens=True,
        )[0].strip()

        entities = parse_entities(raw_output)
        spans = find_entity_spans(sentence, entities)

        highlighted_sentence = create_highlighted_output(
            sentence,
            spans,
        )

        entity_table = [
            [
                span["text"],
                span["type"],
            ]
            for span in spans
        ]

        structured_output = {
            "sentence": sentence,
            "entities": [
                {
                    "text": span["text"],
                    "type": span["type"],
                    "start": span["start"],
                    "end": span["end"],
                }
                for span in spans
            ],
            "raw_model_output": raw_output,
        }

        if not spans:
            structured_output["message"] = (
                "The model returned entities, but their text could not "
                "be matched in the original sentence."
                if entities
                else "No entities were detected."
            )

        return (
            highlighted_sentence,
            entity_table,
            structured_output,
        )

    except Exception as error:
        return (
            [(sentence, None)],
            [],
            {
                "error": str(error),
            },
        )


# ==========================================================
# Clear
# ==========================================================

def clear_all():
    return "", [], [], None


# ==========================================================
# User Interface
# ==========================================================

with gr.Blocks(
    theme=gr.themes.Soft(),
    title="English Named Entity Recognition",
) as demo:

    gr.Markdown(
        """
# 🏷️ English Named Entity Recognition

Enter an English sentence to detect and highlight named entities.

### Supported Entity Types

- 👤 PERSON
- 🏢 ORGANIZATION
- 📍 LOCATION / GPE
- 📅 DATE / TIME
- 🎉 EVENT
- 📦 PRODUCT
- 💰 MONEY
- 🎨 WORK OF ART
- 🗣️ LANGUAGE
- 🏛️ FAC
"""
    )

    with gr.Row():

        with gr.Column(scale=3):

            sentence_input = gr.Textbox(
                label="📝 Original Sentence",
                placeholder=(
                    "Example: Sundar Pichai visited Google headquarters "
                    "in California on July 15, 2026."
                ),
                lines=8,
            )

            with gr.Row():

                submit_button = gr.Button(
                    "🔍 Extract Entities",
                    variant="primary",
                )

                clear_button = gr.Button(
                    "🗑️ Clear",
                )

        with gr.Column(scale=3):

            highlighted_output = gr.HighlightedText(
                label="🎨 Highlighted Sentence",
                color_map=ENTITY_COLORS,
                show_legend=True,
                show_inline_category=True,
                combine_adjacent=True,
            )

            entity_table = gr.Dataframe(
                headers=[
                    "Entity",
                    "Entity Type",
                ],
                datatype=[
                    "str",
                    "str",
                ],
                label="📋 Extracted Entities",
                interactive=False,
            )

    with gr.Accordion(
        "View JSON Output",
        open=False,
    ):
        json_output = gr.JSON(
            label="Structured Output",
        )

    submit_button.click(
        fn=predict,
        inputs=sentence_input,
        outputs=[
            highlighted_output,
            entity_table,
            json_output,
        ],
    )

    sentence_input.submit(
        fn=predict,
        inputs=sentence_input,
        outputs=[
            highlighted_output,
            entity_table,
            json_output,
        ],
    )

    clear_button.click(
        fn=clear_all,
        inputs=[],
        outputs=[
            sentence_input,
            highlighted_output,
            entity_table,
            json_output,
        ],
        queue=False,
    )

    gr.Examples(
        examples=[
            [
                "Sundar Pichai is the CEO of Google and lives in California."
            ],
            [
                "Apple launched the iPhone in September 2025."
            ],
            [
                "Barack Obama visited Paris on January 10, 2024."
            ],
            [
                "Microsoft invested 10 billion dollars in OpenAI."
            ],
            [
                "The FIFA World Cup was held in Qatar in 2022."
            ],
        ],
        inputs=sentence_input,
        label="📚 Try an Example",
    )

    gr.Markdown(
        """
---

### 🚀 About

This application extracts named entities 
"""
    )


demo.queue().launch()