import datetime
import json
from pathlib import Path
from typing import Iterable, Optional, Dict, Any

from ..chunkers.schema import CodeChunk
from ..metadata.chunk_stats import compute_dataset_stats


def _sanitize_text(text: Optional[str]) -> Optional[str]:
    if text is None:
        return None
    return text.replace("\x00", "")


def _infer_language(file_path: str, existing: Optional[str]) -> str:
    if existing:
        return existing
    suffix = Path(file_path).suffix.lower()
    return {
        ".py": "python",
        ".md": "markdown",
        ".txt": "text",
        ".json": "json",
        ".yaml": "yaml",
        ".yml": "yaml",
    }.get(suffix, "text")


def export_chunks_jsonl(
    chunks: Iterable[CodeChunk],
    output_path: Path,
    print_stats: bool = False,
    source_name: Optional[str] = None,
    source_type: str = "local_directory",
) -> Dict[str, Any]:  # ✅ CHANGED: Always returns Dict, not Optional[Dict]
    """
    Export chunks to JSONL with enhanced statistics.
    
    IMPORTANT: ALWAYS returns statistics dictionary, even if not printed.
    
    Args:
        chunks: Iterable of CodeChunk objects
        output_path: Path to output JSONL file
        print_stats: Whether to PRINT statistics to console (stats are ALWAYS computed)
        source_name: Name of the source (optional)
        source_type: Type of source (local_directory, git_repository, etc.)
        
    Returns:
        Enhanced statistics dictionary (ALWAYS returned)
    """
    chunks_list = list(chunks)

    # Ensure output directory exists
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Export chunks to JSONL
    with output_path.open("w", encoding="utf-8") as f:
        for chunk in chunks_list:
            if isinstance(chunk, dict):
                # Handle dict chunks (backward compatibility)
                f.write(json.dumps(chunk, ensure_ascii=False) + "\n")
                continue
            
            # Convert CodeChunk to dict
            code_text = _sanitize_text(chunk.code)
            
            chunk_dict = {
                "chunk_id": chunk.chunk_id,
                "file_path": chunk.file_path,
                "language": _infer_language(chunk.file_path, chunk.language),
                "chunk_type": chunk.chunk_type,
                "code": code_text,
                "ast": (
                    {
                        "symbol_type": chunk.ast.symbol_type,
                        "name": chunk.ast.name,
                        "parent": chunk.ast.parent,
                        "docstring": _sanitize_text(chunk.ast.docstring),
                        "decorators": chunk.ast.decorators,
                        "imports": chunk.ast.imports,
                    }
                    if chunk.ast
                    else None
                ),
                "span": (
                    {
                        "start_byte": chunk.span.start_byte,
                        "end_byte": chunk.span.end_byte,
                        "start_line": chunk.span.start_line,
                        "end_line": chunk.span.end_line,
                    }
                    if chunk.span
                    else None
                ),
                "metadata": chunk.metadata,
                "hierarchy": (
                    {
                        "parent_id": chunk.hierarchy.parent_id,
                        "children_ids": chunk.hierarchy.children_ids,
                        "depth": chunk.hierarchy.depth,
                        "is_primary": chunk.hierarchy.is_primary,
                        "is_extracted": chunk.hierarchy.is_extracted,
                    }
                    if getattr(chunk, "hierarchy", None)
                    else None
                ),
            }

            f.write(json.dumps(chunk_dict, ensure_ascii=False) + "\n")

    # ✅ CHANGED: ALWAYS compute statistics (not conditional)
    code_chunks = [c for c in chunks_list if isinstance(c, CodeChunk)]
    basic_stats = compute_dataset_stats(code_chunks)
    
    # Determine source name if not provided
    if not source_name and chunks_list:
        first_chunk = chunks_list[0]
        source_name = Path(first_chunk.file_path).parent.name
    
    # Build enhanced statistics (ALWAYS)
    enhanced_stats = {
        "source_name": source_name or "unknown",
        "source_type": source_type,
        "processing_timestamp": datetime.datetime.now().isoformat(),
        "total_chunks": basic_stats.get("total_chunks", 0),
        "chunk_type_distribution": basic_stats.get("chunk_type_distribution", {}),
        "language_distribution": basic_stats.get("language_distribution", {}),
        "ast_symbol_distribution": basic_stats.get("ast_symbol_distribution", {}),
        "docstring_coverage_ratio": basic_stats.get("docstring_coverage_ratio", 0),
        "average_code_length_chars": basic_stats.get("average_code_length_chars", 0),
    }
    
    # Only print if requested (stats are ALWAYS computed)
    if print_stats:
        print("\n" + "=" * 60)
        print("📊 DATASET STATISTICS")
        print("=" * 60)
        print(f"📁 Source: {enhanced_stats['source_name']} ({enhanced_stats['source_type']})")
        print(f"🕒 Processed: {enhanced_stats['processing_timestamp']}")
        print(f"🧩 Total Chunks: {enhanced_stats['total_chunks']:,}")
        print(f"📄 Files: {len({c.file_path for c in code_chunks})}")
        
        print(f"\n📊 Chunk Type Distribution:")
        for chunk_type, count in enhanced_stats["chunk_type_distribution"].items():
            percentage = (count / enhanced_stats["total_chunks"]) * 100
            print(f"  {chunk_type:12s}: {count:6d} ({percentage:5.1f}%)")
        
        print(f"\n🌐 Language Distribution:")
        for language, count in enhanced_stats["language_distribution"].items():
            percentage = (count / enhanced_stats["total_chunks"]) * 100
            print(f"  {language:12s}: {count:6d} ({percentage:5.1f}%)")
        
        print(f"\n✅ Docstring Coverage: {enhanced_stats['docstring_coverage_ratio'] * 100:.1f}%")
        print(f"📏 Avg Chunk Size: {enhanced_stats['average_code_length_chars']:.0f} chars")
        print("=" * 60)
    
    # ✅ CHANGED: ALWAYS return statistics (no more None)
    return enhanced_stats


# Example usage (unchanged)
if __name__ == "__main__":
    from pathlib import Path
    from ..chunkers.ast_parser import extract_ast_chunks

    example_file = Path("path/to/example.py")
    chunks = extract_ast_chunks(example_file)
    export_chunks_jsonl(chunks, Path("output/example_chunks.jsonl"), print_stats=True)

# import datetime
# import json
# from pathlib import Path
# from typing import Iterable, Optional, Dict, Any

# from ..chunkers.schema import CodeChunk
# from ..metadata.chunk_stats import compute_dataset_stats


# def _sanitize_text(text: Optional[str]) -> Optional[str]:
#     if text is None:
#         return None
#     return text.replace("\x00", "")


# def _infer_language(file_path: str, existing: Optional[str]) -> str:
#     if existing:
#         return existing
#     suffix = Path(file_path).suffix.lower()
#     return {
#         ".py": "python",
#         ".md": "markdown",
#         ".txt": "text",
#         ".json": "json",
#         ".yaml": "yaml",
#         ".yml": "yaml",
#     }.get(suffix, "text")


# def export_chunks_jsonl(
#     chunks: Iterable[CodeChunk],
#     output_path: Path,
#     print_stats: bool = False,
#     source_name: Optional[str] = None,
#     source_type: str = "local_directory",
# ) -> Optional[Dict[str, Any]]:
#     """
#     Export chunks to JSONL with enhanced statistics.
    
#     Args:
#         chunks: Iterable of CodeChunk objects
#         output_path: Path to output JSONL file
#         print_stats: Whether to print statistics
#         source_name: Name of the source (optional)
#         source_type: Type of source (local_directory, git_repository, etc.)
        
#     Returns:
#         Enhanced statistics dictionary if print_stats=True, None otherwise
#     """
#     chunks_list = list(chunks)

#     # Ensure output directory exists
#     output_path.parent.mkdir(parents=True, exist_ok=True)

#     # Export chunks to JSONL
#     with output_path.open("w", encoding="utf-8") as f:
#         for chunk in chunks_list:
#             if isinstance(chunk, dict):
#                 # Handle dict chunks (backward compatibility)
#                 f.write(json.dumps(chunk, ensure_ascii=False) + "\n")
#                 continue
            
#             # Convert CodeChunk to dict
#             code_text = _sanitize_text(chunk.code)
            
#             chunk_dict = {
#                 "chunk_id": chunk.chunk_id,
#                 "file_path": chunk.file_path,
#                 "language": _infer_language(chunk.file_path, chunk.language),
#                 "chunk_type": chunk.chunk_type,
#                 "code": code_text,
#                 "ast": (
#                     {
#                         "symbol_type": chunk.ast.symbol_type,
#                         "name": chunk.ast.name,
#                         "parent": chunk.ast.parent,
#                         "docstring": _sanitize_text(chunk.ast.docstring),
#                         "decorators": chunk.ast.decorators,
#                         "imports": chunk.ast.imports,
#                     }
#                     if chunk.ast
#                     else None
#                 ),
#                 "span": (
#                     {
#                         "start_byte": chunk.span.start_byte,
#                         "end_byte": chunk.span.end_byte,
#                         "start_line": chunk.span.start_line,
#                         "end_line": chunk.span.end_line,
#                     }
#                     if chunk.span
#                     else None
#                 ),
#                 "metadata": chunk.metadata,
#                 "hierarchy": (
#                     {
#                         "parent_id": chunk.hierarchy.parent_id,
#                         "children_ids": chunk.hierarchy.children_ids,
#                         "depth": chunk.hierarchy.depth,
#                         "is_primary": chunk.hierarchy.is_primary,
#                         "is_extracted": chunk.hierarchy.is_extracted,
#                     }
#                     if getattr(chunk, "hierarchy", None)
#                     else None
#                 ),
#             }

#             f.write(json.dumps(chunk_dict, ensure_ascii=False) + "\n")

#     # Generate and return enhanced statistics if requested
#     if print_stats:
#         # Filter to only CodeChunk objects
#         code_chunks = [c for c in chunks_list if isinstance(c, CodeChunk)]
        
#         # Compute basic statistics
#         basic_stats = compute_dataset_stats(code_chunks)
        
#         # Determine source name if not provided
#         if not source_name and chunks_list:
#             first_chunk = chunks_list[0]
#             source_name = Path(first_chunk.file_path).parent.name
        
#         # Create enhanced statistics - FLAT structure (not nested)
#         enhanced_stats = {
#             # Source metadata (matches repo_exporter format)
#             "source_name": source_name or "unknown",
#             "source_type": source_type,
#             "processing_timestamp": datetime.datetime.now().isoformat(),
            
#             # Original statistics (flat, not nested)
#             "total_chunks": basic_stats.get("total_chunks", 0),
#             "chunk_type_distribution": basic_stats.get("chunk_type_distribution", {}),
#             "language_distribution": basic_stats.get("language_distribution", {}),
#             "ast_symbol_distribution": basic_stats.get("ast_symbol_distribution", {}),
#             "docstring_coverage_ratio": basic_stats.get("docstring_coverage_ratio", 0),
#             "average_code_length_chars": basic_stats.get("average_code_length_chars", 0),
            
#             # Optional: keep original_stats for backward compatibility
#             "_original_stats": basic_stats,  # Underscore indicates internal/backup
#         }
        
#         # Print statistics in user-friendly format
#         print("\n" + "=" * 60)
#         print("📊 DATASET STATISTICS")
#         print("=" * 60)
#         print(f"📁 Source: {enhanced_stats['source_name']} ({enhanced_stats['source_type']})")
#         print(f"🕒 Processed: {enhanced_stats['processing_timestamp']}")
#         print(f"🧩 Total Chunks: {enhanced_stats['total_chunks']:,}")
#         print(f"📄 Files: {len({c.file_path for c in code_chunks})}")
        
#         # Chunk type distribution
#         print("\n📊 Chunk Type Distribution:")
#         for chunk_type, count in enhanced_stats["chunk_type_distribution"].items():
#             percentage = (count / enhanced_stats["total_chunks"]) * 100
#             print(f"  {chunk_type:12s}: {count:6d} ({percentage:5.1f}%)")
        
#         # Language distribution
#         print("\n🌐 Language Distribution:")
#         for language, count in enhanced_stats["language_distribution"].items():
#             percentage = (count / enhanced_stats["total_chunks"]) * 100
#             print(f"  {language:12s}: {count:6d} ({percentage:5.1f}%)")
        
#         # Quality metrics
#         print(f"\n✅ Docstring Coverage: {enhanced_stats['docstring_coverage_ratio'] * 100:.1f}%")
#         print(f"📏 Avg Chunk Size: {enhanced_stats['average_code_length_chars']:.0f} chars")
        
#         return enhanced_stats
    
#     return None


# # Example usage (unchanged)
# if __name__ == "__main__":
#     from pathlib import Path
#     from ..chunkers.ast_parser import extract_ast_chunks

#     example_file = Path("path/to/example.py")
#     chunks = extract_ast_chunks(example_file)
#     export_chunks_jsonl(chunks, Path("output/example_chunks.jsonl"), print_stats=True)

# import datetime
# import json
# from pathlib import Path
# from typing import Iterable, Optional, Dict, Any

# from ..chunkers.schema import CodeChunk
# from ..metadata.chunk_stats import compute_dataset_stats


# def _sanitize_text(text: Optional[str]) -> Optional[str]:
#     if text is None:
#         return None
#     return text.replace("\x00", "")


# def _infer_language(file_path: str, existing: Optional[str]) -> str:
#     if existing:
#         return existing
#     suffix = Path(file_path).suffix.lower()
#     return {
#         ".py": "python",
#         ".md": "markdown",
#         ".txt": "text",
#         ".json": "json",
#         ".yaml": "yaml",
#         ".yml": "yaml",
#     }.get(suffix, "text")


# def export_chunks_jsonl(
#     chunks: Iterable[CodeChunk],
#     output_path: Path,
#     print_stats: bool = False,
#     source_name: Optional[str] = None,
#     source_type: str = "local_directory",
# ) -> Optional[Dict[str, Any]]:
#     """
#     Export chunks to JSONL with enhanced statistics.
    
#     Args:
#         chunks: Iterable of CodeChunk objects
#         output_path: Path to output JSONL file
#         print_stats: Whether to print statistics
#         source_name: Name of the source (optional)
#         source_type: Type of source (local_directory, git_repository, etc.)
        
#     Returns:
#         Enhanced statistics dictionary if print_stats=True, None otherwise
#     """
#     chunks_list = list(chunks)

#     # Ensure output directory exists
#     output_path.parent.mkdir(parents=True, exist_ok=True)

#     # Export chunks to JSONL
#     with output_path.open("w", encoding="utf-8") as f:
#         for chunk in chunks_list:
#             if isinstance(chunk, dict):
#                 # Handle dict chunks (backward compatibility)
#                 f.write(json.dumps(chunk, ensure_ascii=False) + "\n")
#                 continue
            
#             # Convert CodeChunk to dict
#             code_text = _sanitize_text(chunk.code)
            
#             chunk_dict = {
#                 "chunk_id": chunk.chunk_id,
#                 "file_path": chunk.file_path,
#                 "language": _infer_language(chunk.file_path, chunk.language),
#                 "chunk_type": chunk.chunk_type,
#                 "code": code_text,
#                 "ast": (
#                     {
#                         "symbol_type": chunk.ast.symbol_type,
#                         "name": chunk.ast.name,
#                         "parent": chunk.ast.parent,
#                         "docstring": _sanitize_text(chunk.ast.docstring),
#                         "decorators": chunk.ast.decorators,
#                         "imports": chunk.ast.imports,
#                     }
#                     if chunk.ast
#                     else None
#                 ),
#                 "span": (
#                     {
#                         "start_byte": chunk.span.start_byte,
#                         "end_byte": chunk.span.end_byte,
#                         "start_line": chunk.span.start_line,
#                         "end_line": chunk.span.end_line,
#                     }
#                     if chunk.span
#                     else None
#                 ),
#                 "metadata": chunk.metadata,
#                 "hierarchy": (
#                     {
#                         "parent_id": chunk.hierarchy.parent_id,
#                         "children_ids": chunk.hierarchy.children_ids,
#                         "depth": chunk.hierarchy.depth,
#                         "is_primary": chunk.hierarchy.is_primary,
#                         "is_extracted": chunk.hierarchy.is_extracted,
#                     }
#                     if getattr(chunk, "hierarchy", None)
#                     else None
#                 ),
#             }

#             f.write(json.dumps(chunk_dict, ensure_ascii=False) + "\n")

#     # Generate and return enhanced statistics if requested
#     if print_stats:
#         # Filter to only CodeChunk objects
#         code_chunks = [c for c in chunks_list if isinstance(c, CodeChunk)]
        
#         # Compute basic statistics
#         basic_stats = compute_dataset_stats(code_chunks)
        
#         # Determine source name if not provided
#         if not source_name and chunks_list:
#             first_chunk = chunks_list[0]
#             source_name = Path(first_chunk.file_path).parent.name
        
#         # Create enhanced statistics - FLAT structure (not nested)
#         enhanced_stats = {
#             # Source metadata (matches repo_exporter format)
#             "source_name": source_name or "unknown",
#             "source_type": source_type,
#             "processing_timestamp": datetime.datetime.now().isoformat(),
            
#             # Original statistics (flat, not nested)
#             "total_chunks": basic_stats.get("total_chunks", 0),
#             "chunk_type_distribution": basic_stats.get("chunk_type_distribution", {}),
#             "language_distribution": basic_stats.get("language_distribution", {}),
#             "ast_symbol_distribution": basic_stats.get("ast_symbol_distribution", {}),
#             "docstring_coverage_ratio": basic_stats.get("docstring_coverage_ratio", 0),
#             "average_code_length_chars": basic_stats.get("average_code_length_chars", 0),
            
#             # Optional: keep original_stats for backward compatibility
#             "_original_stats": basic_stats,  # Underscore indicates internal/backup
#         }
        
#         # Print statistics in user-friendly format
#         print("\n" + "=" * 60)
#         print("📊 DATASET STATISTICS")
#         print("=" * 60)
#         print(f"📁 Source: {enhanced_stats['source_name']} ({enhanced_stats['source_type']})")
#         print(f"🕒 Processed: {enhanced_stats['processing_timestamp']}")
#         print(f"🧩 Total Chunks: {enhanced_stats['total_chunks']:,}")
#         print(f"📄 Files: {len({c.file_path for c in code_chunks})}")
        
#         # Chunk type distribution
#         print("\n📊 Chunk Type Distribution:")
#         for chunk_type, count in enhanced_stats["chunk_type_distribution"].items():
#             percentage = (count / enhanced_stats["total_chunks"]) * 100
#             print(f"  {chunk_type:12s}: {count:6d} ({percentage:5.1f}%)")
        
#         # Language distribution
#         print("\n🌐 Language Distribution:")
#         for language, count in enhanced_stats["language_distribution"].items():
#             percentage = (count / enhanced_stats["total_chunks"]) * 100
#             print(f"  {language:12s}: {count:6d} ({percentage:5.1f}%)")
        
#         # Quality metrics
#         print(f"\n✅ Docstring Coverage: {enhanced_stats['docstring_coverage_ratio'] * 100:.1f}%")
#         print(f"📏 Avg Chunk Size: {enhanced_stats['average_code_length_chars']:.0f} chars")
        
#         return enhanced_stats
    
#     return None


# # Example usage (unchanged)
# if __name__ == "__main__":
#     from pathlib import Path
#     from ..chunkers.ast_parser import extract_ast_chunks

#     example_file = Path("path/to/example.py")
#     chunks = extract_ast_chunks(example_file)
#     export_chunks_jsonl(chunks, Path("output/example_chunks.jsonl"), print_stats=True)

# import datetime
# import json
# from pathlib import Path
# from typing import Iterable, Optional , Dict, Any

# from ..chunkers.schema import CodeChunk
# from ..metadata.chunk_stats import compute_dataset_stats


# def _sanitize_text(text: Optional[str]) -> Optional[str]:
#     if text is None:
#         return None
#     return text.replace("\x00", "")


# def _infer_language(file_path: str, existing: Optional[str]) -> str:
#     if existing:
#         return existing
#     suffix = Path(file_path).suffix.lower()
#     return {
#         ".py": "python",
#         ".md": "markdown",
#         ".txt": "text",
#         ".json": "json",
#         ".yaml": "yaml",
#         ".yml": "yaml",
#     }.get(suffix, "text")


# # def export_chunks_jsonl(
# #     chunks: Iterable[CodeChunk],
# #     output_path: Path,
# #     print_stats: bool = False,
# # ) -> None:
# def export_chunks_jsonl(
#     chunks: Iterable[CodeChunk],
#     output_path: Path,
#     print_stats: bool = False,
#     source_name: Optional[str] = None,
#     source_type: str = "local_directory",
# ) -> Optional[Dict[str, Any]]:  # Return enhanced stats

#     chunks_list = list(chunks)

#     output_path.parent.mkdir(parents=True, exist_ok=True)

#     with output_path.open("w", encoding="utf-8") as f:
#         for chunk in chunks_list:
#             if isinstance(chunk, dict):
#                 f.write(json.dumps(chunk, ensure_ascii=False) + "\n")
#                 continue
#             code_text = _sanitize_text(chunk.code)

#             chunk_dict = {
#                 "chunk_id": chunk.chunk_id,
#                 "file_path": chunk.file_path,
#                 "language": _infer_language(chunk.file_path, chunk.language),
#                 "chunk_type": chunk.chunk_type,
#                 "code": code_text,
#                 "ast": (
#                     {
#                         "symbol_type": chunk.ast.symbol_type,
#                         "name": chunk.ast.name,
#                         "parent": chunk.ast.parent,
#                         "docstring": _sanitize_text(chunk.ast.docstring),
#                         "decorators": chunk.ast.decorators,
#                         "imports": chunk.ast.imports,
#                     }
#                     if chunk.ast
#                     else None
#                 ),
#                 "span": (
#                     {
#                         "start_byte": chunk.span.start_byte,
#                         "end_byte": chunk.span.end_byte,
#                         "start_line": chunk.span.start_line,
#                         "end_line": chunk.span.end_line,
#                     }
#                     if chunk.span
#                     else None
#                 ),
#                 "metadata": chunk.metadata,
#                 "hierarchy": (
#                     {
#                         "parent_id": chunk.hierarchy.parent_id,
#                         "children_ids": chunk.hierarchy.children_ids,
#                         "depth": chunk.hierarchy.depth,
#                         "is_primary": chunk.hierarchy.is_primary,
#                         "is_extracted": chunk.hierarchy.is_extracted,
#                     }
#                     if getattr(chunk, "hierarchy", None)
#                     else None
#                 ),
#             }

#             f.write(json.dumps(chunk_dict, ensure_ascii=False) + "\n")

#     if print_stats:
#         code_chunks = [c for c in chunks_list if isinstance(c, CodeChunk)]
#         stats = compute_dataset_stats(code_chunks)
        
#         # Determine source name if not provided
#         if not source_name and chunks_list:
#             first_chunk = chunks_list[0]
#             source_name = Path(first_chunk.file_path).parent.name
        
#         # Enhanced statistics for professional output
#         enhanced_stats = {
#             "source_name": source_name or "unknown",
#             "source_type": source_type,
#             "processing_timestamp": datetime.datetime.now().isoformat(),
#             "original_stats": stats,  # Keep original structure
#         }
        
#         print("\n=== Dataset Statistics ===")
#         for k, v in enhanced_stats.items():
#             if k == "original_stats":
#                 print(f"Chunk Statistics:")
#                 for sk, sv in v.items():
#                     print(f"  {sk}: {sv}")
#             else:
#                 print(f"{k}: {v}")
        
#         return enhanced_stats
    
#     return None

# # Example usage
# if __name__ == "__main__":
#     from pathlib import Path
#     from ..chunkers.ast_parser import extract_ast_chunks

#     example_file = Path("path/to/example.py")
#     chunks = extract_ast_chunks(example_file)
#     export_chunks_jsonl(chunks, Path("output/example_chunks.jsonl"), print_stats=True)     
