"""Tests for entity/relationship extraction with mocked LLM."""

from __future__ import annotations

import json
from unittest.mock import AsyncMock

import pytest

from kb_arena.graph.extractor import _validate_result, extract_document
from kb_arena.models.graph import ExtractionResult

AWS_CORPUS = "aws-compute"

_VALID_LLM_RESPONSE = {
    "entities": [
        {
            "id": "aws.lambda.invoke",
            "name": "InvokeFunction",
            "fqn": "aws.lambda.invoke",
            "type": "Process",
            "description": "Invokes an AWS Lambda function.",
            "properties": {},
            "aliases": [],
        },
        {
            "id": "aws.iam.execution-role",
            "name": "ExecutionRole",
            "fqn": "aws.iam.execution-role",
            "type": "Constraint",
            "description": "IAM role that Lambda assumes for execution.",
            "properties": {},
            "aliases": [],
        },
    ],
    "relationships": [
        {
            "source_fqn": "aws.lambda.invoke",
            "target_fqn": "aws.iam.execution-role",
            "type": "DEPENDS_ON",
            "properties": {},
        }
    ],
}


def test_validate_result_accepts_valid_types():
    result = _validate_result(_VALID_LLM_RESPONSE, AWS_CORPUS, "lambda-invoke")
    assert len(result.entities) == 2
    assert len(result.relationships) == 1
    assert result.entities[0].fqn == "aws.lambda.invoke"
    assert result.relationships[0].type == "DEPENDS_ON"


def test_validate_result_rejects_unknown_node_type():
    bad = {
        "entities": [
            {
                "id": "x",
                "name": "X",
                "fqn": "x",
                "type": "UnknownNode",
                "description": "",
                "properties": {},
                "aliases": [],
            }
        ],
        "relationships": [],
    }
    result = _validate_result(bad, AWS_CORPUS, "s1")
    assert result.entities == []


def test_validate_result_rejects_unknown_rel_type():
    bad = {
        "entities": list(_VALID_LLM_RESPONSE["entities"]),
        "relationships": [
            {
                "source_fqn": "aws.lambda.invoke",
                "target_fqn": "aws.iam.execution-role",
                "type": "INVENTED_REL",
                "properties": {},
            }
        ],
    }
    result = _validate_result(bad, AWS_CORPUS, "s1")
    assert result.relationships == []


def test_validate_result_keeps_cross_section_relationships():
    """As of v0.6.0, cross-section relationships are NOT dropped at validation time.

    Per-section validation no longer requires both endpoints to live in the same
    extraction batch. The global FQN union check happens later in run_extraction
    (after every section has been extracted), so multi-hop graph queries
    structurally work.
    """
    payload = {
        "entities": [_VALID_LLM_RESPONSE["entities"][0]],  # only aws.lambda.invoke
        "relationships": list(_VALID_LLM_RESPONSE["relationships"]),  # refs aws.iam.execution-role
    }
    result = _validate_result(payload, AWS_CORPUS, "s1")
    # Edge with both endpoints set is kept; the global pass deduplicates later.
    assert len(result.relationships) == len(_VALID_LLM_RESPONSE["relationships"])


@pytest.mark.asyncio
async def test_extract_document_calls_llm_per_section(sample_document):
    mock_llm = AsyncMock()
    from kb_arena.llm.client import LLMResponse

    mock_llm.extract.return_value = LLMResponse(text=json.dumps(_VALID_LLM_RESPONSE))

    from kb_arena.graph.extractor import _build_system_prompt

    result = await extract_document(sample_document, mock_llm, _build_system_prompt(AWS_CORPUS))

    # Called once per section in sample_document
    assert mock_llm.extract.call_count == len(sample_document.sections)
    assert isinstance(result, ExtractionResult)


@pytest.mark.asyncio
async def test_extract_document_stamps_source_doc_id(sample_document):
    """Entities must carry their source document id so graph retrieval can map
    back to section-level ground truth (graph IR fix)."""
    mock_llm = AsyncMock()
    from kb_arena.llm.client import LLMResponse

    mock_llm.extract.return_value = LLMResponse(text=json.dumps(_VALID_LLM_RESPONSE))

    from kb_arena.graph.extractor import _build_system_prompt

    result = await extract_document(sample_document, mock_llm, _build_system_prompt(AWS_CORPUS))

    assert result.entities
    assert all(e.source_doc_id == sample_document.id for e in result.entities)
    # source_section_id is still set per-section
    assert all(e.source_section_id for e in result.entities)


@pytest.mark.asyncio
async def test_extract_document_handles_bad_json(sample_document):
    mock_llm = AsyncMock()
    from kb_arena.llm.client import LLMResponse

    mock_llm.extract.return_value = LLMResponse(text="not json at all")

    from kb_arena.graph.extractor import _build_system_prompt

    result = await extract_document(sample_document, mock_llm, _build_system_prompt(AWS_CORPUS))
    # Should not crash — returns empty result
    assert isinstance(result, ExtractionResult)
