"""Span attributes demo: entities carry per-span attribute groups (sentiment,
assertion, role), single- or multi-label, scoped with applies_to.

Two examples: a clinical note (conditions with assertion status, medications,
procedures) and a mixed product review (products with sentiment, people with
role). Free-typed text falls back to the review schema.

Response contract:
  {"spans": [{"start", "end", "text", "label", "score",
              "attributes": {"<group>": [{"label": str, "score": float}]}}]}
"""

import json
import os

MOCK = os.environ.get("GLINER_MOCK", "1") == "1"
MODEL_ID = os.environ.get("MODEL_ID", "fastino/gliner2.5-multi-v1")

DEMO = {
    "title": "Span attribute extraction",
    "subtitle": "Extract descriptive labels alongside spans for richer context.",
    "model_name": "GLiNER 2.5",
    "model_variant": "0.3B parameters, f32, CPU",
    "examples": [
        {"chip": "Clinical note", "fixture": "clinical.json",
         "code": (
             'from gliner2 import AutoExtractor, AttributeGroup\n'
             '\n'
             'extractor = AutoExtractor.from_pretrained("fastino/gliner2.5-multi-v1", map_location="cpu")\n'
             'extractor.float()\n'
             '\n'
             'note = open("chart_note.txt").read()\n'
             '\n'
             'schema = (extractor.create_schema()\n'
             '    .entities({"medication": "A medication or drug mentioned in the note",\n'
             '               "condition": "A medical condition, diagnosis, or symptom",\n'
             '               "procedure": "A medical procedure or test"})\n'
             '    .entity_attributes({\n'
             '        "assertion": AttributeGroup(\n'
             '            labels=["present", "absent", "possible", "historical"],\n'
             '            applies_to=["condition"])}))\n'
             '\n'
             'result = extractor.extract(note, schema, include_spans=True, include_confidence=True)\n'
             '# -> {"entities": {"condition": [{"text": "chest pain", "assertion": {"label": "present"}, ...}], ...}}'
         )},
        {"chip": "Mixed review", "fixture": "review.json",
         "code": (
             'from gliner2 import AutoExtractor, AttributeGroup\n'
             '\n'
             'extractor = AutoExtractor.from_pretrained("fastino/gliner2.5-multi-v1", map_location="cpu")\n'
             'extractor.float()\n'
             '\n'
             'review = open("product_review.txt").read()\n'
             '\n'
             'schema = (extractor.create_schema()\n'
             '    .entities({"feature_mention": "A feature, component, or aspect of a product mentioned in the review"})\n'
             '    .entity_attributes({\n'
             '        "sentiment": AttributeGroup(\n'
             '            labels=["positive", "neutral", "negative"],\n'
             '            applies_to=["feature_mention"])}))\n'
             '\n'
             'result = extractor.extract(review, schema, include_spans=True, include_confidence=True)\n'
             '# -> {"entities": {"feature_mention": [{"text": "keyboard", "sentiment": {"label": "positive"}, ...},\n'
             '#                            {"text": "docking station", "sentiment": {"label": "negative"}, ...}], ...}}'
         )},
    ],
}

# Per-example schemas: entity types + attribute groups differ by domain.
# Keys match the fixture .txt basenames.
EXAMPLE_SCHEMAS = {
    "clinical": {
        "entities": {
            "medication": "A medication or drug mentioned in the note",
            "condition": "A medical condition, diagnosis, or symptom",
            "procedure": "A medical procedure or test",
        },
        "attributes": {
            "assertion": (["present", "absent", "possible", "historical"], ["condition"]),
        },
    },
    "review": {
        "threshold": 0.15,  # features in complaint clauses score below the 0.5 default
        "entities": {
            "feature_mention": "A feature, component, or aspect of a product mentioned in the review",
        },
        "attributes": {
            "sentiment": (["positive", "neutral", "negative"], ["feature_mention"]),
        },
    },
}

_HERE = os.path.dirname(os.path.abspath(__file__))


def config():
    """DEMO for /api/config, with result previews baked from fixtures."""
    out = {k: v for k, v in DEMO.items() if k != "code"}
    # per-example code lives on each example dict
    out["mock"] = MOCK
    out["examples"] = []
    for ex in DEMO["examples"]:
        ex2 = {**ex, "text": _text(ex)}
        ex2["result"] = _result_preview(ex["fixture"])
        out["examples"].append(ex2)
    return out


def warmup():
    if MOCK:
        for ex in DEMO["examples"]:
            _fixture(ex["fixture"], _text(ex))
    else:
        _load_real()


def infer(data):
    text = (data.get("text") or "").replace("\r", "")
    if not text.strip():
        return {"spans": []}
    if MOCK:
        return _infer_mock(text)
    return _infer_real(text)


def _text(ex):
    txt = ex.get("text")
    if txt is None:
        with open(os.path.join(_HERE, "fixtures", ex["fixture"].replace(".json", ".txt"))) as f:
            txt = f.read().rstrip("\n")
    return txt


def _schema_key(text):
    for ex in DEMO["examples"]:
        if _text(ex) == text:
            return ex["fixture"].replace(".json", "")
    return None


def _result_preview(fixture_name, per_type=2):
    """Notebook-style Out[1] preview: the API response shape, truncated per label."""
    with open(os.path.join(_HERE, "fixtures", fixture_name)) as f:
        out = json.load(f)
    entities = {}
    for sp in out["spans"]:
        entities.setdefault(sp["label"], [])
        if len(entities[sp["label"]]) < per_type:
            entities[sp["label"]].append(sp)
    counts = {}
    for sp in out["spans"]:
        counts[sp["label"]] = counts.get(sp["label"], 0) + 1
    body = []
    for label, entries in entities.items():
        entry_strs = []
        for e in entries:
            attrs = {g: [{"label": a["label"], "score": a["score"]} for a in es]
                     for g, es in e.get("attributes", {}).items() if es}
            base = (f'{{"text": {json.dumps(e["text"])}, "confidence": {e["score"]}, '
                    f'"start": {e["start"]}, "end": {e["end"]}')
            if attrs:
                base += f', "attributes": {json.dumps(attrs)}'
            base += "}"
            entry_strs.append(base)
        more = counts[label] - len(entries)
        parts = ", ".join(entry_strs)
        if more > 0:
            parts += f",\n        # ... {more} more {label} span{'' if more == 1 else 's'}"
        body.append(f'        {json.dumps(label)}: [\n            {parts}\n        ]')
    return f'{{\n    "entities": {{\n' + ",\n".join(body) + "\n    }\n}"


def _fixture(name, expect_text=None):
    with open(os.path.join(_HERE, "fixtures", name)) as f:
        out = json.load(f)
    if expect_text is not None:
        for sp in out.get("spans", []):
            assert expect_text[sp["start"]:sp["end"]] == sp["text"], \
                f"{name}: bad offsets for {sp['text']!r}"
    return out


def _infer_mock(text):
    for ex in DEMO["examples"]:
        if _text(ex) == text:
            return _fixture(ex["fixture"])
    return {"spans": []}


_model = None


def _load_real():
    global _model
    from gliner2 import AutoExtractor
    _model = AutoExtractor.from_pretrained(MODEL_ID, map_location="cpu")
    _model.float()


def _infer_real(text):
    from gliner2 import AttributeGroup

    key = _schema_key(text)
    cfg = EXAMPLE_SCHEMAS.get(key, EXAMPLE_SCHEMAS["review"])

    schema = _model.create_schema().entities(cfg["entities"])
    attr_groups = {}
    for group, (labels, applies_to) in cfg["attributes"].items():
        kwargs = {"applies_to": applies_to} if applies_to else {}
        attr_groups[group] = AttributeGroup(labels=labels, **kwargs)
    if attr_groups:
        schema = schema.entity_attributes(attr_groups)

    result = _model.extract(text, schema, threshold=cfg.get("threshold", 0.5),
                            include_spans=True, include_confidence=True)

    spans = []
    for label, entries in result.get("entities", {}).items():
        for e in entries:
            span = {
                "start": e["start"], "end": e["end"], "text": e["text"],
                "label": label, "score": e.get("confidence", 1.0),
                "attributes": {},
            }
            for k, v in e.items():
                if k not in ("text", "start", "end", "confidence"):
                    if isinstance(v, dict) and "label" in v:
                        span["attributes"][k] = [{"label": v["label"], "score": v.get("confidence", 1.0)}]
                    elif isinstance(v, list):
                        span["attributes"][k] = [{"label": x["label"], "score": x.get("confidence", 1.0)} for x in v]
            spans.append(span)
    spans.sort(key=lambda s: s["start"])
    return {"spans": spans}