"""Pydantic schemas for the GraphRAG plugin."""

from __future__ import annotations

from datetime import datetime
from typing import Any

from pydantic import BaseModel, Field


class RelationshipCreate(BaseModel):
    source_type: str = Field(..., max_length=50, description="Entity type of the source (e.g. 'contact', 'company')")
    source_id: str = Field(..., description="UUID of the source entity")
    target_type: str = Field(..., max_length=50, description="Entity type of the target (e.g. 'email', 'task')")
    target_id: str = Field(..., description="UUID of the target entity")
    relationship_type: str = Field(..., max_length=50, description="Type of relationship (e.g. 'works_for', 'has_email')")
    metadata: dict[str, Any] | None = Field(None, description="Arbitrary metadata about the relationship")


class RelationshipRead(BaseModel):
    id: str
    source_type: str
    source_id: str
    target_type: str
    target_id: str
    relationship_type: str
    metadata: dict[str, Any] | None = None
    owner_id: str | None = None
    created_at: datetime | None = None


class TraverseRequest(BaseModel):
    source_type: str = Field(..., max_length=50, description="Entity type of the starting node")
    source_id: str = Field(..., description="UUID of the starting node")
    max_hops: int = Field(3, ge=1, le=10, description="Maximum traversal depth")
    relationship_types: list[str] | None = Field(None, description="Optional filter by relationship types")


class GraphNode(BaseModel):
    entity_type: str
    entity_id: str
    depth: int
    path: list[str] = []


class GraphEdge(BaseModel):
    source_type: str
    source_id: str
    target_type: str
    target_id: str
    relationship_type: str
    metadata: dict[str, Any] | None = None


class TraverseResponse(BaseModel):
    nodes: list[GraphNode]
    edges: list[GraphEdge]
    total_nodes: int
    total_edges: int