from __future__ import annotations

from ont_platform.core.extraction.lightweight_extractor import ExtractionResult
from ont_platform.storage.candidate_repository import CandidateRepository
from ont_platform.storage.models import Base, CandidateSource, ReviewStatus
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker


def _session() -> Session:
    engine = create_engine("sqlite:///:memory:")
    Base.metadata.create_all(engine)
    return sessionmaker(bind=engine)()


def test_repository_saves_lightweight_candidates_with_evidence() -> None:
    db = _session()
    repository = CandidateRepository(db)
    result = ExtractionResult(
        entities=[
            {
                "id": "E_alice",
                "label": "Alice",
                "type": "person",
                "confidence": 0.91,
                "evidence_ids": ["EV_1"],
            }
        ],
        relations=[],
        evidence_spans=[
            {
                "id": "EV_1",
                "text": "Alice works at Acme.",
                "start_offset": 0,
                "end_offset": 20,
            }
        ],
        warnings=[],
    )

    batch = repository.save_lightweight_result(
        project_id="proj_1",
        document_id="doc_1",
        result=result,
        source_trust=0.8,
        validation_passed=True,
    )

    assert len(batch.entities) == 1
    entity = batch.entities[0]
    assert entity.source_type == CandidateSource.LIGHTWEIGHT
    assert entity.review_status == ReviewStatus.PENDING
    assert entity.evidence_ids == ["EV_1"]
    assert repository.candidate_has_valid_evidence(entity)


def test_repository_saves_ontocast_candidates_on_separate_source_path() -> None:
    db = _session()
    repository = CandidateRepository(db)
    batch = repository.save_ontocast_result(
        project_id="proj_1",
        document_id="doc_1",
        result={
            "entities": [
                {
                    "id": "E_graph",
                    "label": "GraphUpdate",
                    "entity_type": "concept",
                    "confidence": 0.72,
                    "evidence_ids": ["EV_graph"],
                }
            ],
            "relations": [],
            "evidence_spans": [
                {
                    "id": "EV_graph",
                    "text": "OntoCast proposed a graph update.",
                    "start_offset": 0,
                    "end_offset": 34,
                }
            ],
        },
        source_trust=0.7,
        validation_passed=False,
    )

    entity = batch.entities[0]
    assert entity.source_type == CandidateSource.ONTOCAST
    assert entity.created_by == "ontocast"
    assert entity.validation_passed is False
    assert entity.metadata_["source_type"] == "ontocast"