"""Tests para extract_relations_mrebel con stubs de modelo."""

import os
import sys

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

from python.functions.datascience.extract_relations_mrebel import extract_relations_mrebel
from python.types.datascience.entity_candidate import EntityCandidate
from python.types.datascience.relation_candidate import RelationCandidate


# ---------------------------------------------------------------------------
# Stubs
# ---------------------------------------------------------------------------

class _TokenizerStub:
    """Tokenizer stub que devuelve inputs triviales y decodifica el wire format canonico."""

    def __init__(self, decoded_output: str = ""):
        self._decoded = decoded_output

    def __call__(self, text, return_tensors=None, max_length=512, truncation=True):
        return {"input_ids": [[1, 2, 3]]}

    def decode(self, token_ids, skip_special_tokens=True):
        return self._decoded


class _ModelStub:
    """Modelo stub que devuelve tokens triviales."""

    def generate(self, input_ids=None, num_beams=4, length_penalty=1.0, max_length=256, **kwargs):
        return [[10, 11, 12]]


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

def test_flujo_completo_con_stub_produce_relation_candidates_correctos():
    # Wire format canonico con un triplet valido
    decoded = "<triplet> Pablo Isla <per> Inditex <org> employer"
    tok = _TokenizerStub(decoded_output=decoded)
    model = _ModelStub()

    entities = [
        EntityCandidate(name="Pablo Isla", type_label="PER", confidence=0.95),
        EntityCandidate(name="Inditex", type_label="ORG", confidence=0.92),
    ]
    text = "Pablo Isla es el presidente de Inditex."

    result = extract_relations_mrebel(text, entities, tok, model)

    assert len(result) == 1
    rc = result[0]
    assert isinstance(rc, RelationCandidate)
    assert rc.from_name == "Pablo Isla"
    assert rc.to_name == "Inditex"
    assert rc.relation_type == "employer"
    assert rc.confidence == 1.0


def test_menos_de_2_entidades_retorna_vacio():
    tok = _TokenizerStub()
    model = _ModelStub()
    entities = [EntityCandidate(name="Pablo Isla", type_label="PER")]
    result = extract_relations_mrebel("Texto cualquiera.", entities, tok, model)
    assert result == []


def test_texto_vacio_retorna_vacio():
    tok = _TokenizerStub()
    model = _ModelStub()
    entities = [
        EntityCandidate(name="A", type_label="PER"),
        EntityCandidate(name="B", type_label="ORG"),
    ]
    assert extract_relations_mrebel("", entities, tok, model) == []


def test_triplets_no_alineables_se_descartan():
    # El stub emite entidades que no estan en la lista
    decoded = "<triplet> Ghost Entity <per> Unknown Org <org> some relation"
    tok = _TokenizerStub(decoded_output=decoded)
    model = _ModelStub()

    entities = [
        EntityCandidate(name="Pablo Isla", type_label="PER"),
        EntityCandidate(name="Inditex", type_label="ORG"),
    ]
    result = extract_relations_mrebel("Texto largo suficiente.", entities, tok, model)
    assert result == []


def test_multiples_frases_generan_multiples_candidates():
    # El stub siempre emite el mismo triplet valido — una por frase
    decoded = "<triplet> Pablo Isla <per> Inditex <org> employer"
    tok = _TokenizerStub(decoded_output=decoded)
    model = _ModelStub()

    entities = [
        EntityCandidate(name="Pablo Isla", type_label="PER"),
        EntityCandidate(name="Inditex", type_label="ORG"),
    ]
    # Dos frases separadas por ". "
    text = "Pablo Isla es el presidente de Inditex. Inditex tiene sedes en todo el mundo."

    result = extract_relations_mrebel(text, entities, tok, model)
    # Puede haber 1 o 2 dependiendo de la dedup — lo importante es que no es vacio
    assert len(result) >= 1
    assert all(isinstance(rc, RelationCandidate) for rc in result)
