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)

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)

# Custom CSS for modern look
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
:root {
  --primary: #4f46e5;
  --secondary: #6366f1;
  --background: #f9fafb;
  --card-bg: #ffffff;
  --text: #1f2937;
  --muted: #6b7280;
}
body {
  background: var(--background) !important;
  font-family: 'Inter', sans-serif;
  color: var(--text) !important;
}
header.brand {
  padding: 2rem 0;
  text-align: center;
}
header.brand .logo {
  font-size: 2rem;
  font-weight: 700;
  color: var(--primary);
}
header.brand .subtitle {
  margin-top: 0.2rem;
  font-size: 0.9rem;
  color: var(--muted);
}
.gradio-container {
  max-width: 800px;
  margin: auto;
  padding: 1rem;
}
.card {
  background: var(--card-bg);
  padding: 1.5rem;
  border-radius: 0.75rem;
  box-shadow: 0 4px 10px rgba(0,0,0,0.05);
  margin-bottom: 1.5rem;
}
.gr-button.primary {
  background: var(--primary) !important;
  color: #fff !important;
  border-radius: 0.5rem;
  padding: 0.6rem 1.2rem;
}
"""

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>  
          <div class=\"subtitle\">Compact • Modern • Screenshot-Ready</div>  
        </header>
        """
    )

    with gr.Tabs():
        # NER Tab
        with gr.TabItem("🔍 Named Entity Recognition"):
            with gr.Row(elem_classes="card", gap="small"):
                with gr.Column(scale=2):
                    txt1 = gr.Textbox(label="Input Text", lines=5, placeholder="Enter text to extract entities...")
                    types1 = gr.Textbox(label="Entity Types (CSV)", value="person, organization, location, date, title, topic")
                    with gr.Accordion("Optional Descriptions", open=False):
                        desc1 = gr.Textbox(lines=3, placeholder="person: Full name\norganization: Companies\n...")
                    btn1 = gr.Button("Extract Entities", variant="primary")
                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", gap="small"):
                with gr.Column(scale=2):
                    txt2 = gr.Textbox(label="Input Text", lines=5, placeholder="Enter text to classify...")
                    task2 = gr.Textbox(label="Task Name", value="sentiment_analysis")
                    labs2 = gr.Textbox(label="Labels (CSV)", value="positive, negative, neutral")
                    with gr.Accordion("Optional Label Descriptions", open=False):
                        desc2 = gr.Textbox(lines=3, placeholder="positive: Positive sentiment\n...")
                    multi2 = gr.Checkbox(label="Multi-label?", value=False)
                    btn2 = gr.Button("Classify Text", variant="primary")
                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", gap="small"):
                with gr.Column(scale=2):
                    txt3 = gr.Textbox(label="Input Text", lines=5, placeholder="Enter text for structure extraction...")
                    struct3 = gr.Code(language="json", label="Schema (JSON)", lines=8, value=json.dumps({
                        "product": [
                            "name::str::Product name and model",
                            "price::str::Product price",
                            "features::list::Key features",
                            "category::str::Product category"
                        ]
                    }, indent=2))
                    btn3 = gr.Button("Extract Structure", variant="primary")
                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)