"""Tests for free relation extraction and relation type normalization.

Covers:
- normalize_rel_type: alias mapping, canonical pass-through, unknown types pass-through,
  sanitization (spaces, hyphens → underscores, lowercase)
- Extraction prompts: verify that prompts no longer restrict to canonical types
- Round-trip: extraction with non-canonical type → normalization → accepted
"""

from __future__ import annotations

import json

from conftest import FakeLLMProvider, run_async

# ===========================================================================
# Tests: normalize_rel_type
# ===========================================================================


class TestNormalizeRelType:
    """Tests for the normalize_rel_type function."""

    def test_canonical_type_passes_through(self):
        """Canonical types like works_at are returned unchanged."""
        from arandu.constants import normalize_rel_type

        assert normalize_rel_type("works_at") == "works_at"
        assert normalize_rel_type("manages") == "manages"
        assert normalize_rel_type("lives_in") == "lives_in"

    def test_alias_maps_to_canonical(self):
        """Known aliases are mapped to their canonical type."""
        from arandu.constants import normalize_rel_type

        assert normalize_rel_type("boss") == "reports_to"
        assert normalize_rel_type("employee_of") == "works_at"
        assert normalize_rel_type("spouse") == "partner_of"
        assert normalize_rel_type("sibling") == "family_of"
        assert normalize_rel_type("colleague") == "works_with"

    def test_unknown_type_passes_through(self):
        """Types outside canonical set and aliases pass through after sanitization."""
        from arandu.constants import normalize_rel_type

        assert normalize_rel_type("mentored_by") == "mentored_by"
        assert normalize_rel_type("inspired_by") == "inspired_by"
        assert normalize_rel_type("competed_with") == "competed_with"
        assert normalize_rel_type("taught_at") == "taught_at"

    def test_sanitization_lowercase(self):
        """Input is lowercased."""
        from arandu.constants import normalize_rel_type

        assert normalize_rel_type("WORKS_AT") == "works_at"
        assert normalize_rel_type("Mentored_By") == "mentored_by"

    def test_sanitization_spaces_to_underscores(self):
        """Spaces are replaced with underscores."""
        from arandu.constants import normalize_rel_type

        assert normalize_rel_type("works at") == "works_at"
        assert normalize_rel_type("mentored by") == "mentored_by"

    def test_sanitization_hyphens_to_underscores(self):
        """Hyphens are replaced with underscores."""
        from arandu.constants import normalize_rel_type

        assert normalize_rel_type("works-at") == "works_at"
        assert normalize_rel_type("mentored-by") == "mentored_by"

    def test_sanitization_strips_whitespace(self):
        """Leading/trailing whitespace is stripped."""
        from arandu.constants import normalize_rel_type

        assert normalize_rel_type("  works_at  ") == "works_at"
        assert normalize_rel_type("\tmentored_by\n") == "mentored_by"


# ===========================================================================
# Tests: CANONICAL_REL_TYPES is reference, not filter
# ===========================================================================


class TestCanonicalRelTypesAsReference:
    """Tests that CANONICAL_REL_TYPES exists as reference but doesn't filter."""

    def test_canonical_set_exists(self):
        """CANONICAL_REL_TYPES is still available as a set."""
        from arandu.constants import CANONICAL_REL_TYPES

        assert isinstance(CANONICAL_REL_TYPES, set)
        assert "works_at" in CANONICAL_REL_TYPES
        assert len(CANONICAL_REL_TYPES) >= 11

    def test_aliases_still_exist(self):
        """_REL_TYPE_ALIASES is still available for synonym mapping."""
        from arandu.constants import _REL_TYPE_ALIASES

        assert isinstance(_REL_TYPE_ALIASES, dict)
        assert "boss" in _REL_TYPE_ALIASES


# ===========================================================================
# Tests: Extraction prompts allow free types
# ===========================================================================


class TestExtractionPromptsFreetype:
    """Tests that extraction prompts no longer restrict to canonical types."""

    def test_multi_pass_relation_prompt_has_no_allowed_types(self):
        """RELATION_EXTRACTION_PROMPT does not contain 'Allowed types:'."""
        from arandu.write.extract import RELATION_EXTRACTION_PROMPT

        assert "Allowed types:" not in RELATION_EXTRACTION_PROMPT
        assert "snake_case" in RELATION_EXTRACTION_PROMPT

    def test_prompts_list_common_types_as_examples(self):
        """Prompts list common rel_types as illustrative examples."""
        from arandu.write.extract import RELATION_EXTRACTION_PROMPT

        assert "works_at" in RELATION_EXTRACTION_PROMPT
        assert "lives_in" in RELATION_EXTRACTION_PROMPT


# ===========================================================================
# Tests: Free extraction round-trip
# ===========================================================================


class TestFreeExtractionRoundTrip:
    """Tests that non-canonical types are extracted and normalized correctly."""

    def test_extraction_with_non_canonical_type(self):
        """LLM returning a non-canonical rel_type is parsed successfully."""
        from arandu.write.extract import extract_from_message

        response = json.dumps(
            {
                "entities": [
                    {"name": "user", "type": "person"},
                    {"name": "Professor Silva", "type": "person"},
                ],
                "facts": [
                    {"subject": "user", "fact": "User was mentored by Professor Silva", "confidence": 0.85},
                ],
                "relations": [
                    {"source": "user", "target": "Professor Silva", "rel_type": "mentored_by"},
                ],
            }
        )

        llm = FakeLLMProvider([response])
        result, _meta = run_async(extract_from_message("Professor Silva me orientou no mestrado", llm=llm))

        assert len(result.relations) == 1
        assert result.relations[0].rel_type == "mentored_by"

    def test_normalize_after_extraction(self):
        """Non-canonical types from extraction normalize correctly."""
        from arandu.constants import normalize_rel_type

        # Simulating what happens after extraction: the rel_type goes through normalize
        assert normalize_rel_type("mentored_by") == "mentored_by"
        assert normalize_rel_type("inspired by") == "inspired_by"
        assert normalize_rel_type("COMPETED-WITH") == "competed_with"
