"""Tests para extract_graph_gliner2.

Usa un stub GLiNER2 para validar el contrato sin descargar el modelo real.
"""

from __future__ import annotations

import os
import sys

import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..", ".."))

from python.functions.datascience.extract_graph_gliner2 import extract_graph_gliner2


class _Schema:
    def entities(self, labels):
        self._entities = labels
        return self

    def relations(self, labels):
        self._relations = labels
        return self


class _StubModel:
    """Stub que devuelve entidades y relaciones conocidas."""

    _extract_result = {
        "entities": {"person": ["Pablo Isla"], "organization": ["Inditex"]},
        "relation_extraction": {"ceo_of": [("Pablo Isla", "Inditex")]},
    }

    def create_schema(self):
        return _Schema()

    def extract(self, text, schema=None, threshold=0.3, include_confidence=False):
        return self._extract_result


def test_output_tiene_claves_entities_relation_extraction_elapsed_s():
    """output tiene claves entities relation_extraction elapsed_s"""
    result = extract_graph_gliner2(
        text="Pablo Isla es CEO de Inditex.",
        entity_labels=["person", "organization"],
        relation_labels=["ceo_of"],
        model=_StubModel(),
    )
    assert "entities" in result
    assert "relation_extraction" in result
    assert "elapsed_s" in result
    assert isinstance(result["elapsed_s"], float)


def test_stub_model_retorna_shape_correcto():
    """stub model retorna shape correcto"""
    result = extract_graph_gliner2(
        text="Texto cualquiera.",
        entity_labels=["person"],
        relation_labels=["works_at"],
        model=_StubModel(),
        threshold=0.3,
    )
    assert result["entities"] == {"person": ["Pablo Isla"], "organization": ["Inditex"]}
    assert "ceo_of" in result["relation_extraction"]
