import spaces  # MUST come before torch / any CUDA-touching import
import torch
import gradio as gr
import json
import re

from gliner2 import AutoExtractor

MODEL_ID = "fastino/gliner2.5-multi-v1"

model = AutoExtractor.from_pretrained(MODEL_ID, map_location="cuda")
model.eval()

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""


def _parse_labels(labels_text):
    """Parse comma-separated labels into a clean list."""
    if not labels_text or not labels_text.strip():
        return []
    labels = [l.strip() for l in labels_text.split(",") if l.strip()]
    return labels


def _parse_class_schema(schema_text):
    """Parse classification schema from text like: sentiment: positive, negative, neutral"""
    result = {}
    if not schema_text or not schema_text.strip():
        return result
    for line in schema_text.strip().split("\n"):
        if ":" in line:
            task, labels_str = line.split(":", 1)
            task = task.strip()
            labels = [l.strip() for l in labels_str.split(",") if l.strip()]
            if task and labels:
                result[task] = labels
    return result


def _format_json(obj):
    """Pretty-print JSON for display."""
    return json.dumps(obj, indent=2, ensure_ascii=False, default=str)


def _entities_to_highlights(text, result):
    """Turn GLiNER2 entity spans into gr.HighlightedText (text, label) tuples."""
    spans = []
    for label, items in (result.get("entities") or {}).items():
        for item in items or []:
            if not isinstance(item, dict):
                continue
            start, end = item.get("start"), item.get("end")
            if start is None or end is None:
                continue
            spans.append((int(start), int(end), label, float(item.get("confidence") or 0.0)))

    # Resolve overlapping spans: keep the most confident, then the longest.
    spans.sort(key=lambda s: (-s[3], -(s[1] - s[0])))
    kept = []
    for start, end, label, conf in spans:
        if any(start < k_end and end > k_start for k_start, k_end, _, _ in kept):
            continue
        kept.append((start, end, label, conf))
    kept.sort(key=lambda s: s[0])

    highlights = []
    cursor = 0
    for start, end, label, _conf in kept:
        if start > cursor:
            highlights.append((text[cursor:start], None))
        highlights.append((text[start:end], label))
        cursor = end
    if cursor < len(text):
        highlights.append((text[cursor:], None))
    return highlights or [(text, None)]


@spaces.GPU(duration=30)
def extract_entities(text, labels_text):
    """Extract named entities from text using zero-shot GLiNER2.5.

    Returns the input text as a list of (substring, entity_label) pairs, where
    entity_label is null for spans that are not part of an entity.

    Args:
        text: The input text to extract entities from.
        labels_text: Comma-separated entity labels to detect (e.g. "person, organization, location").
    """
    labels = _parse_labels(labels_text)
    if not text.strip():
        return [("Please enter some text.", None)]
    if not labels:
        return [("Please enter at least one entity label.", None)]
    result = model.extract_entities(
        text,
        labels,
        include_confidence=True,
        include_spans=True,
    )
    return _entities_to_highlights(text, result)


@spaces.GPU(duration=30)
def classify_text(text, schema_text):
    """Classify text into categories using zero-shot classification with GLiNER2.5.

    Args:
        text: The input text to classify.
        schema_text: Classification schema, one task per line in format 'task: label1, label2, ...'.
    """
    schema = _parse_class_schema(schema_text)
    if not text.strip():
        return "Please enter some text."
    if not schema:
        return "Please enter a classification schema (e.g. 'sentiment: positive, negative, neutral')."
    result = model.classify_text(text, schema)
    return _format_json(result)


@spaces.GPU(duration=30)
def extract_relations(text, labels_text, include_confidence=True, include_spans=True):
    """Extract relations between entities from text using GLiNER2.5.

    Args:
        text: The input text to extract relations from.
        labels_text: Comma-separated relation labels to detect (e.g. "works_for, located_in").
        include_confidence: Whether to include confidence scores.
        include_spans: Whether to include character spans.
    """
    labels = _parse_labels(labels_text)
    if not text.strip():
        return "Please enter some text."
    if not labels:
        return "Please enter at least one relation label."
    result = model.extract_relations(
        text,
        labels,
        include_confidence=include_confidence,
        include_spans=include_spans,
    )
    return _format_json(result)


@spaces.GPU(duration=30)
def extract_structured(text, schema_text):
    """Extract structured JSON data from text using GLiNER2.5.

    Args:
        text: The input text to extract structured data from.
        schema_text: JSON schema description, one field per line in format 'field: type::description'.
    """
    if not text.strip():
        return "Please enter some text."
    # Parse schema: field_name::type::description (one per line)
    schema = {}
    for line in schema_text.strip().split("\n"):
        line = line.strip()
        if not line:
            continue
        parts = line.split("::", 2)
        if len(parts) >= 1:
            field = parts[0].strip()
            dtype = parts[1].strip() if len(parts) > 1 else "str"
            desc = parts[2].strip() if len(parts) > 2 else ""
            entry = f"{dtype}::{desc}" if desc else dtype
            schema[field] = [entry] if dtype != "list" else [f"list::{desc}" if desc else "list"]
    if not schema:
        return "Please enter a schema (e.g. 'name::str::Product name')."
    result = model.extract_json(text, schema)
    return _format_json(result)


with gr.Blocks(title="GLiNER2.5 Multi — Information Extraction") as demo:
    gr.Markdown("""
    # 🔍 GLiNER2.5 Multi — Zero-Shot Information Extraction

    Multilingual, multi-task information extraction with [fastino/gliner2.5-multi-v1](https://huggingface.co/fastino/gliner2.5-multi-v1) (287M params, mDeBERTa-v3 encoder).

    Define your own labels at inference time — no retraining needed. Supports entity recognition, text classification, relation extraction, and structured data extraction across multiple languages.
    """)

    with gr.Tabs():
        with gr.Tab("🏷️ Entity Extraction"):
            gr.Markdown("Extract named entities with custom labels.")
            with gr.Row():
                with gr.Column():
                    ner_text = gr.Textbox(
                        label="Input Text",
                        placeholder="Enter text to analyze…",
                        lines=5,
                        value="Apple CEO Tim Cook announced iPhone 15 in Cupertino yesterday. The event was held at Apple Park.",
                    )
                    ner_labels = gr.Textbox(
                        label="Entity Labels (comma-separated)",
                        value="company, person, product, location",
                        placeholder="person, organization, location…",
                    )
                    ner_btn = gr.Button("Extract Entities", variant="primary")
                with gr.Column():
                    ner_output = gr.HighlightedText(
                        label="Extracted Entities",
                        combine_adjacent=True,
                        show_legend=True,
                    )

            gr.Examples(
                examples=[
                    ["Apple CEO Tim Cook announced iPhone 15 in Cupertino yesterday. The event was held at Apple Park.", "company, person, product, location"],
                    ["Barcelona defeated Real Madrid 3-1 at Camp Nou. Lewandowski scored twice for Barça.", "team, player, city, stadium"],
                    ["Marie Curie was born in Warsaw and later moved to Paris to work at the Sorbonne.", "person, city, country, organization"],
                ],
                inputs=[ner_text, ner_labels],
                outputs=ner_output,
                fn=extract_entities,
                cache_examples=True,
                cache_mode="lazy",
            )

        with gr.Tab("📋 Text Classification"):
            gr.Markdown("Classify text into custom categories (zero-shot).")
            with gr.Row():
                with gr.Column():
                    cls_text = gr.Textbox(
                        label="Input Text",
                        placeholder="Enter text to classify…",
                        lines=3,
                        value="This laptop has amazing performance but terrible battery life!",
                    )
                    cls_schema = gr.Textbox(
                        label="Classification Schema (one task per line: task: label1, label2, …)",
                        value="sentiment: positive, negative, neutral",
                        lines=3,
                    )
                    cls_btn = gr.Button("Classify Text", variant="primary")
                with gr.Column():
                    cls_output = gr.Code(label="Result (JSON)", language="json", lines=8)

            gr.Examples(
                examples=[
                    ["This laptop has amazing performance but terrible battery life!", "sentiment: positive, negative, neutral"],
                    ["Breaking: Tech giant acquires AI startup for $2B in landmark deal.", "topic: technology, business, politics, sports"],
                    ["Le film était captivant du début à la fin, avec des acteurs brillants.", "sentiment: positif, négatif, neutre"],
                ],
                inputs=[cls_text, cls_schema],
                outputs=cls_output,
                fn=classify_text,
                cache_examples=True,
                cache_mode="lazy",
            )

        with gr.Tab("🔗 Relation Extraction"):
            gr.Markdown("Detect relationships between entities in text.")
            with gr.Row():
                with gr.Column():
                    rel_text = gr.Textbox(
                        label="Input Text",
                        placeholder="Enter text to analyze…",
                        lines=4,
                        value="Alice works for Acme Corp in Paris. Bob joined Acme last year and lives in London.",
                    )
                    rel_labels = gr.Textbox(
                        label="Relation Labels (comma-separated)",
                        value="works_for, located_in",
                        placeholder="works_for, located_in, founded_by…",
                    )
                    with gr.Accordion("Options", open=False):
                        rel_conf = gr.Checkbox(label="Include confidence scores", value=True)
                        rel_spans = gr.Checkbox(label="Include character spans", value=True)
                    rel_btn = gr.Button("Extract Relations", variant="primary")
                with gr.Column():
                    rel_output = gr.Code(label="Result (JSON)", language="json", lines=15)

            gr.Examples(
                examples=[
                    ["Alice works for Acme Corp in Paris. Bob joined Acme last year and lives in London.", "works_for, located_in"],
                    ["Google was founded by Larry Page and Sergey Brin in Mountain View.", "founded_by, located_in"],
                    ["John Smith married Jane Doe in 2015 in New York City.", "married_to, located_in"],
                ],
                inputs=[rel_text, rel_labels],
                outputs=rel_output,
                fn=extract_relations,
                cache_examples=True,
                cache_mode="lazy",
            )

        with gr.Tab("📦 Structured Data Extraction"):
            gr.Markdown("Parse text into structured JSON records with typed fields.")
            with gr.Row():
                with gr.Column():
                    json_text = gr.Textbox(
                        label="Input Text",
                        placeholder="Enter text to extract structured data from…",
                        lines=4,
                        value="iPhone 15 Pro Max with 256GB storage, A17 Pro chip, priced at $1199. Available in titanium and black colors.",
                    )
                    json_schema = gr.Textbox(
                        label="Schema (one field per line: field::type::description)",
                        value="name::str::Full product name and model\nstorage::str::Storage capacity\nprocessor::str::Chip or processor\nprice::str::Product price with currency\ncolors::list::Available color options",
                        lines=5,
                    )
                    json_btn = gr.Button("Extract Structured Data", variant="primary")
                with gr.Column():
                    json_output = gr.Code(label="Result (JSON)", language="json", lines=12)

            gr.Examples(
                examples=[
                    ["iPhone 15 Pro Max with 256GB storage, A17 Pro chip, priced at $1199. Available in titanium and black colors.", "name::str::Full product name and model\nstorage::str::Storage capacity\nprocessor::str::Chip or processor\nprice::str::Product price with currency\ncolors::list::Available color options"],
                    ["Alice bought apples for $3.50 at Whole Foods. Bob purchased oranges for $2.00 at Trader Joe's.", "buyer::str::Name of buyer\nitem::str::Item purchased\nprice::str::Price paid\nstore::str::Store name"],
                ],
                inputs=[json_text, json_schema],
                outputs=json_output,
                fn=extract_structured,
                cache_examples=True,
                cache_mode="lazy",
            )

    # Wire buttons
    ner_btn.click(
        extract_entities,
        inputs=[ner_text, ner_labels],
        outputs=ner_output,
        api_name="extract_entities",
    )
    cls_btn.click(
        classify_text,
        inputs=[cls_text, cls_schema],
        outputs=cls_output,
        api_name="classify_text",
    )
    rel_btn.click(
        extract_relations,
        inputs=[rel_text, rel_labels, rel_conf, rel_spans],
        outputs=rel_output,
        api_name="extract_relations",
    )
    json_btn.click(
        extract_structured,
        inputs=[json_text, json_schema],
        outputs=json_output,
        api_name="extract_structured",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)