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

# Get API key from environment variable
hf_token = os.getenv("HF_TOKEN")

# Authenticate with Hugging Face
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)

# ——— Clean White Theme & Layout ———
custom_css = """
body {
  background: #ffffff !important;
  font-family: 'Helvetica Neue', sans-serif;
  color: #333333;
}
.gradio-container {
  max-width: 600px;
  padding: 0;
  background: #ffffff;
}
header, .logo, .subtitle {
  border: none !important;
  box-shadow: none !important;
}
.gradio-container * {
  box-shadow: none !important;
}
.card {
  background: #ffffff;
  padding: 15px;
}
label {
  color: #444444;
  font-weight: 600;
}
.gr-textbox textarea,
.gr-code,
.gr-dropdown,
.gr-checkbox,
.gr-button {
  background: #ffffff !important;
  box-shadow: none !important;
}
.accordion-button {
  border: none !important;
  box-shadow: none !important;
  font-weight: 500;
}
.gr-button.primary {
  background: #5b8def;
  color: #ffffff;
}
"""

with gr.Blocks(theme=gr.themes.Base(), css=custom_css) as demo:
    # Header
    gr.HTML("""
    <header style="text-align:center; padding:10px 0;">
      <div class="logo" style="font-size:1.8rem; font-weight:700; color:#333333;">🎯 GLiNER2</div>
      <div class="subtitle" style="font-size:0.85rem; color:#777777;">Compact • White Theme • Screenshot-Ready</div>
    </header>
    """)

    with gr.Tabs():
        # Structure Extraction Tab
        with gr.TabItem("Hierarchical Structure Extraction"):
            with gr.Row(elem_classes="card"):
                with gr.Column(scale=2):
                    txt3 = gr.Textbox(
                        label="Input text", lines=3,
                        value=(
                            "The Acme Pro Laptop 15” features an Intel Core i7 processor, 16GB RAM, 512GB SSD, "
                            "and a 15.6-inch 4K display. Priced at $1,499, it offers Wi-Fi 6, Bluetooth 5.2, and "
                            "a backlit keyboard."
                        )
                    )
                    struct3 = gr.Code(
                        language="json", lines=7,
                        label = "Schema",
                        value=json.dumps({
                            "product": [
                                "name::str::Product name and model",
                                "price::str::Product cost",
                                "features::list::Key product features",
                                "category::[electronics|software|hardware]::str"
                            ]
                        }, indent=2)
                    )
                    btn3 = gr.Button("Predict", variant="primary")
                with gr.Column(scale=1):
                    out3 = gr.Code(language="json", lines=8, label="Output")
                btn3.click(run_struct, [txt3, struct3], out3)

        # NER Tab
        with gr.TabItem("Named Entity Recognition"):
            with gr.Row(elem_classes="card"):
                with gr.Column(scale=2):
                    txt1 = gr.Textbox(
                        label="Text", lines=4,
                        value=(
                            "Dr. Alice Smith, Chief Data Scientist at OpenAI, spoke at the AI Summit "
                            "in San Francisco on June 12, 2025, about advancements in large-scale language "
                            "models, ethical AI guidelines, and real-world GPT-4 Turbo applications."
                        )
                    )
                    types1 = gr.Textbox(label="Types (csv)", value="person, title, organization, event, location, date, topic")
                    with gr.Accordion("Descriptions (opt)", open=False):
                        desc1 = gr.Textbox(lines=4, placeholder=(
                            "person: Full names\n"
                            "title: Roles\n"
                            "organization: Companies\n"
                            "event: Conferences\n"
                            "location: Cities\n"
                            "date: Temporal expressions"
                        ))
                    btn1 = gr.Button("Predict", variant="primary")
                with gr.Column(scale=1):
                    out1 = gr.Code(language="json", lines=8)
                btn1.click(run_ner, [txt1, types1, desc1], out1)

        # Classification Tab
        with gr.TabItem("Text Classification"):
            with gr.Row(elem_classes="card"):
                with gr.Column(scale=2):
                    txt2 = gr.Textbox(
                        label="Text", lines=4,
                        value=(
                            "The Q2 2025 financial report shows a 15% revenue increase driven by cloud "
                            "services, offset by a 12% rise in R&D costs. Overall sentiment is cautiously "
                            "optimistic among stakeholders."
                        )
                    )
                    task2 = gr.Textbox(label="Task", value="financial_sentiment")
                    labs2 = gr.Textbox(label="Labels (csv)", value="positive, negative, neutral, mixed, uncertain")
                    with gr.Accordion("Label Descriptions (opt)", open=False):
                        desc2 = gr.Textbox(lines=3, placeholder=(
                            "positive: Favorable outcomes\n"
                            "negative: Concerns raised\n"
                            "neutral: Balanced reporting"
                        ))
                    multi2 = gr.Checkbox(label="Multi-label?", value=True)
                    btn2 = gr.Button("Predict", variant="primary")
                with gr.Column(scale=1):
                    out2 = gr.Code(language="json", lines=8)
                btn2.click(run_class, [txt2, task2, labs2, desc2, multi2], out2)


    demo.launch(share=False)