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)

