"""
ChunkHive Schema Definitions
=============================

Core data structures for semantic code chunking with hierarchical relationships.

This module defines the foundational data models used throughout ChunkHive for
representing parsed code chunks, their Abstract Syntax Tree (AST) information,
hierarchical relationships, and spatial positioning within source files.

Author: ChunkHive
License: Apache 2.0
Version: 0.2.0
"""

from typing import Dict, List, Optional, Literal, Any
from dataclasses import dataclass, field, asdict
from enum import Enum


# =============================================================================
# TYPE DEFINITIONS
# =============================================================================

SemanticChunkType = Literal[
    "module",         # Python module (file-level)
    "class",          # Class definition
    "function",       # Top-level function
    "method",         # Class method
    "async_function", # Async function definition
    "async_method",   # Async class method
    "context",        # Contextual code block
    "documentation",  # Markdown/RST documentation
    "configuration",  # Config files (JSON, YAML, TOML, INI)
    "notebook",       # Jupyter notebook cell
    "script",         # Shell/batch scripts
    "dockerfile",     # Docker configuration
    "typescript",     # TypeScript code (future)
    "javascript",     # JavaScript code (future)
    "text",           # Plain text content
    "imports",        # Import/require statements
    "code",           # Generic code block
    "unknown"         # Unrecognized type
]

ASTSymbolType = Literal[
    "module", "class", "function", "method", "async_function", "async_method",
    "context", "documentation", "configuration", "notebook", "script",
    "dockerfile", "typescript", "javascript", "text", "imports", "code",
    "unknown"
]


# =============================================================================
# HIERARCHICAL RELATIONSHIP METADATA
# =============================================================================

@dataclass
class HierarchyMetadata:
    """
    Captures the hierarchical relationships between code chunks.
    
    This class models parent-child relationships, depth in the hierarchy tree,
    and provides a complete lineage path for each chunk. Essential for 
    maintaining semantic context when chunks are processed independently.
    
    Attributes:
        parent_id: Unique identifier of the parent chunk (None for root-level)
        children_ids: List of child chunk identifiers
        depth: Depth in the hierarchy tree (0 = module level)
        is_primary: Whether this is a primary chunk (vs. extracted fragment)
        is_extracted: Whether chunk was extracted from a larger context
        lineage: Full path from root to this chunk (list of ancestor IDs)
        sibling_index: Position among siblings (0-indexed)
        
    Example:
        >>> hierarchy = HierarchyMetadata(
        ...     parent_id="module_abc123",
        ...     children_ids=["method_def456", "method_ghi789"],
        ...     depth=1,
        ...     lineage=["module_abc123"]
        ... )
    """
    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
    
    def add_child(self, child_id: str) -> None:
        """Add a child chunk ID to this hierarchy node."""
        if child_id not in self.children_ids:
            self.children_ids.append(child_id)
    
    def get_lineage_path(self) -> str:
        """Return human-readable lineage path (e.g., 'module -> class -> method')."""
        return " -> ".join(self.lineage) if self.lineage else "root"


# =============================================================================
# ABSTRACT SYNTAX TREE METADATA
# =============================================================================

@dataclass
class ASTNodeMetadata:
    """
    Abstract Syntax Tree information extracted during parsing.
    
    Contains semantic information about the code structure as determined by
    language-specific AST parsing (Python's ast module, Tree-sitter, etc.).
    This preserves the semantic meaning beyond raw text.
    
    Attributes:
        symbol_type: Type of AST symbol (function, class, etc.)
        name: Symbol name as it appears in source code
        parent: Name of parent symbol (for nested definitions)
        docstring: Extracted documentation string (if present)
        decorators: List of decorator names applied to this symbol
        imports: Import statements within this chunk's scope
        node_type: Specific AST node type from parser
        return_type: Function/method return type annotation (if present)
        parameters: Function/method parameter names and types
        
    Example:
        >>> ast_meta = ASTNodeMetadata(
        ...     symbol_type="function",
        ...     name="calculate_embeddings",
        ...     docstring="Generate semantic embeddings for code chunks.",
        ...     decorators=["@lru_cache"],
        ...     parameters=["chunks: List[str]", "model: str"]
        ... )
    """
    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
    return_type: Optional[str] = None
    parameters: List[str] = field(default_factory=list)
    
    def has_docstring(self) -> bool:
        """Check if this node has a non-empty docstring."""
        return bool(self.docstring and self.docstring.strip())
    
    def get_signature(self) -> Optional[str]:
        """Generate function signature string if applicable."""
        if self.symbol_type in ["function", "method", "async_function", "async_method"]:
            params = ", ".join(self.parameters) if self.parameters else ""
            return_part = f" -> {self.return_type}" if self.return_type else ""
            return f"{self.name}({params}){return_part}"
        return None


# =============================================================================
# SOURCE LOCATION METADATA
# =============================================================================

@dataclass
class SourceLocation:
    """
    Precise location information for a chunk within its source file.
    
    Provides byte-level and line-level positioning to enable exact extraction,
    modification, and referencing of code chunks. All positions are 0-indexed
    for byte offsets and 1-indexed for line numbers (matching editor conventions).
    
    Attributes:
        start_byte: Starting byte offset in file (0-indexed)
        end_byte: Ending byte offset in file (exclusive, 0-indexed)
        start_line: Starting line number (1-indexed)
        end_line: Ending line number (inclusive, 1-indexed)
        char_count: Total character count (including whitespace)
        line_count: Total number of lines spanned
        
    Example:
        >>> location = SourceLocation(
        ...     start_byte=1024,
        ...     end_byte=2048,
        ...     start_line=42,
        ...     end_line=67
        ... )
        >>> print(f"Spans {location.line_count} lines")
    """
    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
    
    @property
    def line_count(self) -> int:
        """Calculate number of lines spanned by this chunk."""
        if self.start_line and self.end_line:
            return self.end_line - self.start_line + 1
        return 0
    
    @property
    def byte_length(self) -> int:
        """Calculate byte length of this chunk."""
        if self.start_byte is not None and self.end_byte is not None:
            return self.end_byte - self.start_byte
        return 0
    
    def validate(self) -> bool:
        """Validate that location data is consistent."""
        if self.start_byte is not None and self.end_byte is not None:
            if self.end_byte <= self.start_byte:
                return False
        if self.start_line is not None and self.end_line is not None:
            if self.end_line < self.start_line:
                return False
        return True


# =============================================================================
# MAIN CODE CHUNK DATA STRUCTURE
# =============================================================================

@dataclass
class SemanticChunk:
    """
    Core representation of a semantic code chunk.
    
    A SemanticChunk represents a meaningful unit of code that has been extracted
    using AST-based parsing. It contains the actual code, metadata about its
    structure and meaning, hierarchical relationships, and optional AI-generated
    enhancements.
    
    This is the primary data structure used throughout ChunkHive for representing
    parsed code, passing it through processing pipelines, and storing/retrieving
    from databases or vector stores.
    
    Attributes:
        chunk_id: Unique deterministic identifier for this chunk
        file_path: Relative or absolute path to source file
        language: Programming language (python, javascript, etc.)
        chunk_type: Semantic type of this chunk
        code: Actual source code content
        ast: Abstract Syntax Tree metadata
        span: Source location information
        metadata: Additional metadata (repo info, processing info, etc.)
        hierarchy: Hierarchical relationship information
        embedding: Optional vector embedding for semantic search (768-dim typical)
        embedding_model: Name of model used to generate embedding
        ai_analysis: Optional AI-generated analysis and insights
        
    Example:
        >>> chunk = SemanticChunk(
        ...     chunk_id="func_auth_a1b2c3",
        ...     file_path="src/auth.py",
        ...     language="python",
        ...     chunk_type="function",
        ...     code="def authenticate(user, password):\\n    ...",
        ...     ast=ASTNodeMetadata(name="authenticate", symbol_type="function"),
        ...     span=SourceLocation(start_line=42, end_line=67)
        ... )
    """
    chunk_id: str
    file_path: str
    language: str
    chunk_type: SemanticChunkType
    code: str
    ast: ASTNodeMetadata
    span: SourceLocation
    metadata: Dict[str, Any] = field(default_factory=dict)
    hierarchy: HierarchyMetadata = field(default_factory=HierarchyMetadata)
    
    # ✨ NEW: Enhanced fields for AI features
    embedding: Optional[List[float]] = None
    embedding_model: Optional[str] = None
    ai_analysis: Optional[Dict[str, Any]] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """
        Convert chunk to dictionary representation for serialization.
        
        Returns:
            Dictionary containing all chunk data with nested objects converted
        """
        return {
            "chunk_id": self.chunk_id,
            "file_path": self.file_path,
            "language": self.language,
            "chunk_type": self.chunk_type,
            "code": self.code,
            "ast": asdict(self.ast),
            "span": asdict(self.span),
            "metadata": self.metadata,
            "hierarchy": asdict(self.hierarchy),
            "embedding": self.embedding,
            "embedding_model": self.embedding_model,
            "ai_analysis": self.ai_analysis,
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "SemanticChunk":
        """
        Create SemanticChunk from dictionary representation.
        
        Args:
            data: Dictionary containing chunk data
            
        Returns:
            Reconstructed SemanticChunk instance
        """
        return cls(
            chunk_id=data["chunk_id"],
            file_path=data["file_path"],
            language=data["language"],
            chunk_type=data["chunk_type"],
            code=data["code"],
            ast=ASTNodeMetadata(**data["ast"]),
            span=SourceLocation(**data["span"]),
            metadata=data.get("metadata", {}),
            hierarchy=HierarchyMetadata(**data.get("hierarchy", {})),
            embedding=data.get("embedding"),
            embedding_model=data.get("embedding_model"),
            ai_analysis=data.get("ai_analysis"),
        )
    
    def get_code_preview(self, max_lines: int = 5) -> str:
        """
        Get a preview of the code (first N lines).
        
        Args:
            max_lines: Maximum number of lines to include
            
        Returns:
            Preview string with ellipsis if truncated
        """
        lines = self.code.split('\n')
        if len(lines) <= max_lines:
            return self.code
        preview_lines = lines[:max_lines]
        return '\n'.join(preview_lines) + '\n...'
    
    def has_embedding(self) -> bool:
        """Check if this chunk has an embedding vector."""
        return self.embedding is not None and len(self.embedding) > 0
    
    def has_ai_analysis(self) -> bool:
        """Check if this chunk has AI-generated analysis."""
        return self.ai_analysis is not None and len(self.ai_analysis) > 0
    
    def get_full_name(self) -> str:
        """
        Get the fully qualified name of this chunk.
        
        Returns:
            Dotted path like "module.Class.method"
        """
        if self.ast.parent and self.ast.name:
            return f"{self.ast.parent}.{self.ast.name}"
        return self.ast.name or self.chunk_id


# =============================================================================
# BACKWARD COMPATIBILITY ALIASES
# =============================================================================

# Keep old names for backward compatibility with existing code
ChunkType = SemanticChunkType
ChunkHierarchy = HierarchyMetadata
ChunkAST = ASTNodeMetadata
ChunkSpan = SourceLocation
CodeChunk = SemanticChunk


# =============================================================================
# MODULE EXPORTS
# =============================================================================

__all__ = [
    # New names (preferred)
    "SemanticChunkType",
    "ASTSymbolType",
    "HierarchyMetadata",
    "ASTNodeMetadata",
    "SourceLocation",
    "SemanticChunk",
    # Backward compatibility aliases
    "ChunkType",
    "ChunkHierarchy",
    "ChunkAST",
    "ChunkSpan",
    "CodeChunk",
]

# 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)