from typing import List, Optional, Dict, Any
from pydantic import BaseModel, Field
from enum import Enum

class NodeType(str, Enum):
    FUNCTION = "Function"
    CLASS = "Class"
    FILE = "File"
    VARIABLE = "Variable"
    NAMESPACE = "Namespace"

class EdgeType(str, Enum):
    CALLS = "CALLS"
    CONTAINS = "CONTAINS"
    EXTENDS = "EXTENDS"
    IMPLEMENTS = "IMPLEMENTS"
    HAS_PARAMETER = "HAS_PARAMETER"
    RETURNS = "RETURNS"

class GraphNode(BaseModel):
    id: str
    type: NodeType
    name: str
    signature: Optional[str] = None
    file_path: Optional[str] = None
    start_line: Optional[int] = None
    end_line: Optional[int] = None
    summary: Optional[str] = None
    logic_flow: Optional[str] = None
    embedding: Optional[List[float]] = None
    raw_attributes: Dict[str, Any] = Field(default_factory=dict)

class GraphEdge(BaseModel):
    source_id: str
    target_id: str
    type: EdgeType
    properties: Dict[str, Any] = Field(default_factory=dict)