import gradio as gr
from gliner import GLiNER
from huggingface_hub import snapshot_download
from typing import Dict, Union

# ---------------------------------------------------------
# MODEL LOADING
# ---------------------------------------------------------
try:
    model_path = snapshot_download(repo_id="lucasorrentino/ContractNer")
    model = GLiNER.from_pretrained(model_path)
    print("Success: Loaded ContractNer model from HuggingFace Hub.")
except Exception as e:
    print(f"Remote load failed: {e}")
    print("Fallback: Loading base model (urchade/gliner_multi-v2.1).")
    model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")

# ---------------------------------------------------------
# EXAMPLES
# ---------------------------------------------------------
examples = [
    [
        "This EMPLOYMENT AGREEMENT is made effective as of January 1, 2026, by and between Agile Lab S.r.l. (the 'Company') and John Doe (the 'Executive'), serving as Chief Technology Officer.",
        "Parties, EffectiveDate, Role, Organization",
        0.5,
        False,
    ],
    [
        "The Company agrees to pay the Executive an annual base salary of $225,000.00, subject to standard deductions. In addition, the Executive shall receive 5,000 Shares of common stock.",
        "Salary, Shares, Role, Parties, Currency",
        0.3,
        False,
    ],
    [
        "This Lease shall terminate on December 31, 2030 (the 'Termination Date'). The monthly Rent shall be $2,500.00 payable to Smith Properties LLC at 123 Main St, New York.",
        "TerminationDate, Rent, Parties, Address, Location",
        0.4,
        False,
    ],
    [
        "Any dispute arising out of this contract shall be settled by arbitration in the State of Texas in accordance with the Commercial Arbitration Rules.",
        "Court, Regulation, Location",
        0.3,
        False,
    ],
]

# ---------------------------------------------------------
# INFERENCE
# ---------------------------------------------------------
def ner(text: str, labels: str, threshold: float, nested_ner: bool) -> Dict:
    labels_list = [l.strip() for l in labels.split(",")]
    entities = model.predict_entities(
        text, labels_list, flat_ner=not nested_ner, threshold=threshold
    )
    return {
        "text": text,
        "entities": [
            {
                "entity": entity["label"],
                "word": entity["text"],
                "start": entity["start"],
                "end": entity["end"],
                "score": entity["score"],
            }
            for entity in entities
        ],
    }

# ---------------------------------------------------------
# GRADIO UI
# ---------------------------------------------------------
with gr.Blocks(title="ContractNER – Legal Entity Extraction") as demo:

    gr.Markdown(
        """
        # ContractNER – Legal Entity Extraction
        A fine-tuned **GLiNER** model specialized for **Legal Contracts**.
        Extract entities such as Parties, Dates, Salaries, Roles, Regulations, and Financial Terms.

        > **Zero-shot / Open Vocabulary** – define your own labels at inference time.
        """
    )

    with gr.Accordion("How to run this model locally", open=False):
        gr.Markdown("### Installation")
        gr.Code("pip install gliner", language="shell")
        gr.Markdown("### Usage")
        gr.Code(
            '''from gliner import GLiNER

model = GLiNER.from_pretrained("lucasorrentino/ContractNer")

text = """
This EMPLOYMENT AGREEMENT is made effective as of January 1, 2026,
by and between Tech Solutions Inc. ("Company") and Jane Doe ("Executive").
The Executive shall serve as Chief Technology Officer with a salary of $250,000.00.
"""

labels = ["Parties", "EffectiveDate", "Role", "Salary"]
entities = model.predict_entities(text, labels, threshold=0.5)

for entity in entities:
    print(f"{entity[\'text\']} => {entity[\'label\']} ({entity[\'score\']:.2f})")
''',
            language="python",
        )

    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(
                value=examples[0][0],
                label="Contract Text",
                placeholder="Paste your contract snippet here...",
                lines=6,
            )
            labels_input = gr.Textbox(
                value=examples[0][1],
                label="Entity Labels (comma separated)",
                placeholder="Parties, Date, Salary...",
                info="Define which entities you want to extract.",
            )
            with gr.Row():
                threshold = gr.Slider(
                    0, 1, value=0.5, step=0.01,
                    label="Threshold",
                    info="Minimum confidence score to keep a prediction.",
                )
                nested_ner = gr.Checkbox(
                    value=False,
                    label="Nested NER",
                    info="Allow overlapping entities?",
                )
            submit_btn = gr.Button("Extract Entities", variant="primary")

        with gr.Column(scale=2):
            output = gr.HighlightedText(label="Predicted Entities")

    gr.Markdown("### Try these examples:")
    gr.Examples(
        examples,
        fn=ner,
        inputs=[input_text, labels_input, threshold, nested_ner],
        outputs=output,
        cache_examples=True,
    )

    submit_btn.click(fn=ner, inputs=[input_text, labels_input, threshold, nested_ner], outputs=output)
    input_text.submit(fn=ner, inputs=[input_text, labels_input, threshold, nested_ner], outputs=output)
    labels_input.submit(fn=ner, inputs=[input_text, labels_input, threshold, nested_ner], outputs=output)

if __name__ == "__main__":
    demo.queue()
    demo.launch()