import gradio as gr
import json
from gliner2 import GLiNER2
from huggingface_hub import login
import os

# Authenticate with Hugging Face
hf_token = os.getenv("HF_TOKEN")
login(hf_token)

# Load model once
model = GLiNER2.from_pretrained("fastino/gliner2-base-0207")

def run_ner(text, types_csv, descs):
    types = [t.strip() for t in types_csv.split(",") if t.strip()]
    desc_map = {k: v for line in descs.split("\n") if ":" in line for k,v in [line.split(":",1)]}
    inp = desc_map if desc_map else types
    res = model.extract_entities(text=text, entity_types=inp, include_confidence=True)
    return model.pretty_print_results(res, include_confidence=True)


def run_class(text, task, labels_csv, descs, multi):
    labels = [l.strip() for l in labels_csv.split(",") if l.strip()]
    desc_map = {k: v for line in descs.split("\n") if ":" in line for k,v in [line.split(":",1)]}
    inp = desc_map if desc_map else labels
    tasks = {
        task: {
            "labels": list(inp.keys()) if isinstance(inp,dict) else inp,
            "multi_label": multi,
            **({"label_descriptions": inp} if isinstance(inp,dict) else {})
        }
    }
    res = model.classify_text(text=text, tasks=tasks, include_confidence=True)
    return model.pretty_print_results(res, include_confidence=True)


def run_struct(text, struct_json):
    try:
        cfg = json.loads(struct_json)
    except json.JSONDecodeError as e:
        return f"❌ Invalid JSON: {e}"
    res = model.extract_json(text=text, structures=cfg, include_confidence=True)
    return model.pretty_print_results(res, include_confidence=True)

# Simplified CSS - uses default backgrounds
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');

body {
  font-family: 'Inter', sans-serif;
}

header.brand {
  padding: 2rem 0;
  text-align: center;
}

header.brand .logo {
  font-size: 2rem;
  font-weight: 700;
  color: #4f46e5;
}

header.brand .subtitle {
  margin-top: 0.2rem;
  font-size: 0.9rem;
  color: #6b7280;
}

.gr-button.primary {
  background: #4f46e5 !important;
  color: #fff !important;
  border-radius: 0.5rem;
  padding: 0.6rem 1.2rem;
}
"""

# Pre-made examples for each task (5 per tab)
ner_examples = [
    [
        "Barack Obama visited Berlin in July 2013.",
        "person,location,date",
        "person: Full name\nlocation: City\ndate: Month and year"
    ],
    [
        "Apple released the iPhone 13 on September 14, 2021.",
        "organization,product,date",
        "organization: Company name\nproduct: Device name\ndate: Full date"
    ],
    [
        "Elon Musk announced Tesla's new Roadster at the LA Auto Show.",
        "person,organization,event,location",
        "person: Full name\norganization: Company name\nevent: Conference or show\nlocation: Venue"
    ],
    [
        "The UEFA Champions League Final takes place in Istanbul this year.",
        "event,location,date",
        "event: Sports event\nlocation: City\ndate: Year"
    ],
    [
        "Microsoft acquired GitHub in 2018 for $7.5 billion.",
        "organization,organization,date,price",
        "organization: Company name\ndate: Year\nprice: Acquisition value"
    ]
]

class_examples = [
    [
        "The movie was a thrilling experience with stunning visuals.",
        "sentiment",
        "positive,negative,neutral",
        "positive: Positive sentiment\nnegative: Negative sentiment\nneutral: Mixed or neutral",
        False
    ],
    [
        "Our Q1 results were disappointing, with sales down 10%.",
        "financial_sentiment",
        "positive,negative,neutral",
        "positive: Gains\nnegative: Losses\nneutral: Flat",
        False
    ],
    [
        "I love the new interface but dislike the slow loading time.",
        "feedback",
        "praise,complaint,suggestion",
        "praise: Positive feedback\ncomplaint: Negative feedback\nsuggestion: Improvement ideas",
        True
    ],
    [
        "The product meets expectations but could use more features.",
        "review",
        "positive,negative",
        "positive: Meets expectations\nnegative: Lacking",
        False
    ],
    [
        "Customer support was helpful, though response times were slow.",
        "support_sentiment",
        "positive,negative,neutral",
        "positive: Helpful support\nnegative: Unhelpful support\nneutral: Mixed experiences",
        True
    ]
]

struct_examples = [
    [
        "The iPad Pro comes with an M1 chip, 8GB RAM, 256GB storage, and a 12.9-inch display.",
        json.dumps({
            "device": [
                "name::str::Model name",
                "specs::list::Hardware specifications",
                "price::str::Device cost"
            ]
        }, indent=2)
    ],
    [
        "Plan: Write report (Due: May 10), Review code (Due: May 15), Deploy (Due: May 20)",
        json.dumps({
            "tasks": [
                "title::str::Task title",
                "due_date::str::Due date"
            ]
        }, indent=2)
    ],
    [
        "Product: Coffee Mug; Price: $12; Features: ceramic, dishwasher-safe, 12oz capacity.",
        json.dumps({
            "product": [
                "name::str::Product name",
                "price::str::Product price",
                "features::list::Product features"
            ]
        }, indent=2)
    ],
    [
        "Event: AI Conference; Date: August 22, 2025; Location: Paris; Topics: ML, Ethics, Robotics.",
        json.dumps({
            "event": [
                "name::str::Event name",
                "date::str::Event date",
                "location::str::Event location",
                "topics::list::Covered topics"
            ]
        }, indent=2)
    ],
    [
        "Recipe: Pancakes; Ingredients: flour, eggs, milk; Steps: mix, cook, serve.",
        json.dumps({
            "recipe": [
                "title::str::Recipe title",
                "ingredients::list::List of ingredients",
                "steps::list::Preparation steps"
            ]
        }, indent=2)
    ]
]

with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="blue"), css=custom_css) as demo:
    # Header
    gr.HTML(
        """
        <header class=\"brand\">  
          <div class=\"logo\">✨ GLiNER2</div>  
        </header>
        """
    )

    with gr.Tabs():
        # NER Tab
        with gr.TabItem("🔍 Named Entity Recognition"):
            with gr.Row(elem_classes="card"):
                with gr.Column(scale=2):
                    txt1 = gr.Textbox(label="Input Text", lines=5)
                    types1 = gr.Textbox(label="Entity Types (CSV)")
                    with gr.Accordion("Optional Descriptions", open=False):
                        desc1 = gr.Textbox(lines=3)
                    btn1 = gr.Button("Extract Entities", variant="primary")
                    gr.Examples(examples=ner_examples, inputs=[txt1, types1, desc1], outputs=None, fn=lambda *args: None, cache_examples=False)
                with gr.Column(scale=1):
                    out1 = gr.Code(language="json", label="Results", lines=8)
                btn1.click(run_ner, inputs=[txt1, types1, desc1], outputs=out1)

        # Classification Tab
        with gr.TabItem("📝 Text Classification"):
            with gr.Row(elem_classes="card"):
                with gr.Column(scale=2):
                    txt2 = gr.Textbox(label="Input Text", lines=5)
                    task2 = gr.Textbox(label="Task Name")
                    labs2 = gr.Textbox(label="Labels (CSV)")
                    with gr.Accordion("Optional Label Descriptions", open=False):
                        desc2 = gr.Textbox(lines=3)
                    multi2 = gr.Checkbox(label="Multi-label?")
                    btn2 = gr.Button("Classify Text", variant="primary")
                    gr.Examples(examples=class_examples, inputs=[txt2, task2, labs2, desc2, multi2], outputs=None, fn=lambda *args: None, cache_examples=False)
                with gr.Column(scale=1):
                    out2 = gr.Code(language="json", label="Results", lines=8)
                btn2.click(run_class, inputs=[txt2, task2, labs2, desc2, multi2], outputs=out2)

        # Structure Extraction Tab
        with gr.TabItem("📐 Structure Extraction"):
            with gr.Row(elem_classes="card"):
                with gr.Column(scale=2):
                    txt3 = gr.Textbox(label="Input Text", lines=5)
                    struct3 = gr.Code(language="json", label="Schema (JSON)", lines=8)
                    btn3 = gr.Button("Extract Structure", variant="primary")
                    gr.Examples(examples=struct_examples, inputs=[txt3, struct3], outputs=None, fn=lambda *args: None, cache_examples=False)
                with gr.Column(scale=1):
                    out3 = gr.Code(language="json", label="Results", lines=8)
                btn3.click(run_struct, inputs=[txt3, struct3], outputs=out3)

    demo.launch(share=False, width=800)