import json

import pytest
from pydantic import ValidationError

from app.services.memory_extraction import (
    MemoryExtractionResult,
    build_extraction_prompt,
    parse_extraction_response,
)


def test_parse_memory_extraction_result_with_entities_edges_and_evidence():
    raw = json.dumps(
        {
            "entities": [
                {
                    "mention": "Alice",
                    "canonical_name": "Alice",
                    "labels": ["Person"],
                    "aliases": ["A"],
                    "attributes": {"role": "founder"},
                    "summary": "A founder.",
                    "confidence": 0.92,
                }
            ],
            "edges": [
                {
                    "source_entity_ref": "Alice",
                    "target_entity_ref": "Bob",
                    "relation": "KNOWS",
                    "fact": "Alice knows Bob.",
                    "attributes": {},
                    "valid_at": None,
                    "invalid_at": None,
                    "expired_at": None,
                    "confidence": 0.8,
                    "evidence": ["episode-1:0-15"],
                }
            ],
            "episode_summary": "A relationship statement.",
            "unresolved_mentions": ["Bob"],
        }
    )

    result = parse_extraction_response(raw)
    assert isinstance(result, MemoryExtractionResult)
    assert result.entities[0].canonical_name == "Alice"
    assert result.edges[0].relation == "KNOWS"
    assert result.edges[0].evidence == ["episode-1:0-15"]


def test_parse_memory_extraction_result_rejects_unknown_fields_and_bad_confidence():
    payload = {
        "entities": [],
        "edges": [],
        "episode_summary": "summary",
        "unresolved_mentions": [],
        "secret_prompt": "must not be accepted",
    }
    with pytest.raises(ValidationError):
        parse_extraction_response(json.dumps(payload))

    payload.pop("secret_prompt")
    payload["entities"] = [
        {
            "mention": "Alice",
            "canonical_name": "Alice",
            "labels": [],
            "aliases": [],
            "attributes": {},
            "summary": "",
            "confidence": 1.5,
        }
    ]
    with pytest.raises(ValidationError):
        parse_extraction_response(payload)


def test_extraction_prompt_puts_language_instruction_first_and_bounds_episode_size():
    prompt = build_extraction_prompt(
        language="th",
        ontology={"entity_types": ["Person"], "edge_types": ["KNOWS"]},
        episode_text="Alice knows Bob.",
    )
    assert prompt.startswith("IMPORTANT:")
    assert "Return JSON only" in prompt
    assert "entity_refs" in prompt
    assert "Alice knows Bob." in prompt

    with pytest.raises(ValueError, match="episode_too_large"):
        build_extraction_prompt(language="en", ontology={}, episode_text="x" * 20001)