from typing import Dict, List, Optional, Literal, Union
from dataclasses import dataclass, field


# ✅ ADD 'code' to ChunkType
ChunkType = Literal[
    "module",        # Python module
    "class",         # Python class  
    "function",      # Python function
    "method",        # Python method
    "context",       # General context
    "documentation", # Markdown/RST docs
    "configuration", # Config files (JSON, YAML, TOML)
    "notebook",      # Jupyter notebook
    "script",        # Shell scripts
    "dockerfile",    # Docker files
    "typescript",    # TypeScript files
    "javascript",    # JavaScript files
    "text",          # Plain text
    "imports",       # Import statements
    "code",          # ✅ ADDED: Generic code block
    "unknown"        # Unknown file type
]

# ✅ ADD 'code' to ASTSymbolType
ASTSymbolType = Literal[
    "module", "class", "function", "method", "context",
    "documentation", "configuration", "notebook", "script",
    "dockerfile", "typescript", "javascript", "text", 
    "imports", "code",  # ✅ ADDED
    "unknown"
]


@dataclass  
class ChunkHierarchy:
    """Enhanced hierarchical relationship metadata"""
    parent_id: Optional[str] = None
    children_ids: List[str] = field(default_factory=list)
    depth: int = 0
    is_primary: bool = True
    is_extracted: bool = False
    lineage: List[str] = field(default_factory=list)
    sibling_index: int = 0


@dataclass
class ChunkAST:
    symbol_type: Optional[ASTSymbolType] = None
    name: Optional[str] = None
    parent: Optional[str] = None
    docstring: Optional[str] = None
    decorators: List[str] = field(default_factory=list)
    imports: List[str] = field(default_factory=list)
    node_type: Optional[str] = None


@dataclass
class ChunkSpan:
    start_byte: Optional[int] = None
    end_byte: Optional[int] = None
    start_line: Optional[int] = None
    end_line: Optional[int] = None
    char_count: Optional[int] = None


@dataclass
class CodeChunk:
    chunk_id: str
    file_path: str
    language: str
    chunk_type: ChunkType  # ✅ Now accepts 'code' type
    code: str
    ast: ChunkAST
    span: ChunkSpan
    metadata: Dict = field(default_factory=dict)
    hierarchy: ChunkHierarchy = field(default_factory=ChunkHierarchy)