Source code for lalandre_db_neo4j.models

"""
Neo4j Graph Models
"""

from datetime import datetime
from typing import Any, Optional

from lalandre_core.utils import convert_dates_to_strings
from pydantic import BaseModel, Field


def _empty_graph_items() -> list[dict[str, Any]]:
    return []


def _empty_metadata() -> dict[str, Any]:
    return {}




class ActNode(BaseModel):
    """
    Represents an Act as a node in the knowledge graph
    """

    id: int
    celex: str = Field(
        ...,
        min_length=1,
        max_length=100,
        description="Document identifier (EU CELEX format 3YYYYXNNNN, AMF format AMF-*, or other regulatory codes)",
    )
    title: str
    act_type: str  # Stored as string in Neo4j (enum value from ActType)
    language: str  # Stored as string in Neo4j (enum value from LanguageCode)
    adoption_date: Optional[datetime] = None
    force_date: Optional[datetime] = None
    end_date: Optional[datetime] = None
    sector: Optional[int] = None
    level: Optional[int] = Field(
        default=None, ge=1, le=3, description="Regulatory level: 1=EU, 2=National, 3=Decisions"
    )
    official_journal_reference: Optional[str] = None
    eli: Optional[str] = None
    url_eurlex: Optional[str] = None



    def to_neo4j_properties(self) -> dict[str, Any]:
        """Convert to Neo4j node properties"""
        props = self.model_dump(exclude_none=True)
        return convert_dates_to_strings(props, ["adoption_date", "force_date", "end_date"])






class ActRelationship(BaseModel):
    """
    Represents a relationship between two acts
    """

    source_act_id: int
    target_act_id: int
    relation_type: str = Field(..., description="Type of relationship from RelationType enum")
    effect_date: Optional[datetime] = None
    description: Optional[str] = None
    source_subdivision_id: Optional[int] = None
    target_subdivision_id: Optional[int] = None



    def to_neo4j_properties(self) -> dict[str, Any]:
        """Convert to Neo4j relationship properties"""
        props = self.model_dump(
            exclude={"source_act_id", "target_act_id"},
            exclude_none=True,
        )
        return convert_dates_to_strings(props, ["effect_date"])




    def get_neo4j_type(self) -> str:
        """Get Neo4j relationship type (uppercase)"""
        return self.relation_type.upper()






class EntityNode(BaseModel):
    """
    Represents a named legal entity (org, concept, jurisdiction, topic) in the knowledge graph.
    Linked to Acts via MENTIONS relationships.
    """

    name: str = Field(..., min_length=1, description="Canonical entity name")
    entity_type: str = Field(
        ...,
        description="Entity category: ORGANIZATION, CONCEPT, JURISDICTION, or REGULATORY_TOPIC",
    )
    description: Optional[str] = Field(default=None, description="Short 1-sentence description")



    def to_neo4j_properties(self) -> dict[str, Any]:
        """Return the serializable property map for the entity node."""
        return self.model_dump(exclude_none=True)






class CommunityNode(BaseModel):
    """
    Represents a detected community of Acts in the knowledge graph.
    Created by Louvain community detection; linked to Acts via BELONGS_TO.
    """

    id: int = Field(..., description="Community identifier (0-based, sequential)")
    num_acts: int = 0
    num_relations: int = 0
    relation_types: str = Field(
        default="{}", description="JSON-encoded dict of relation type counts, e.g. '{\"AMENDS\": 42}'"
    )
    central_acts: str = Field(
        default="[]", description="JSON-encoded list of central act dicts [{act_id, celex, title, degree}]"
    )
    summary: str = Field(default="", description="Human-readable community summary")
    modularity: Optional[float] = None
    resolution: Optional[float] = None



    def to_neo4j_properties(self) -> dict[str, Any]:
        """Return the serializable property map for the community node."""
        return self.model_dump(exclude_none=True)






class GraphQueryResult(BaseModel):
    """
    Result from a graph query
    Contains act nodes and relationships that form a subgraph
    """

    nodes: list[dict[str, Any]] = Field(default_factory=_empty_graph_items)
    relationships: list[dict[str, Any]] = Field(default_factory=_empty_graph_items)
    metadata: dict[str, Any] = Field(default_factory=_empty_metadata)