#!/usr/bin/env python3
"""
JSONL to Individual JSON Files Converter
Converts JSONL files to separate JSON files with proper naming and formatting.
"""

import json
import argparse
from pathlib import Path
import logging
from tqdm import tqdm
import re


def setup_logging(verbose: bool = False):
    """Setup logging configuration"""
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(level=level, format="%(asctime)s - %(levelname)s - %(message)s")


def sanitize_filename(filename: str) -> str:
    """
    Sanitize filename to be filesystem-safe

    Args:
        filename: Original filename

    Returns:
        Safe filename for filesystem
    """
    # Replace problematic characters
    safe_name = re.sub(r'[<>:"/\\|?*]', "_", filename)

    # Remove or replace other problematic characters
    safe_name = re.sub(r"[^\w\-_.]", "_", safe_name)

    # Remove multiple consecutive underscores
    safe_name = re.sub(r"_+", "_", safe_name)

    # Trim underscores from start/end
    safe_name = safe_name.strip("_")

    # Ensure it's not empty
    if not safe_name:
        safe_name = "unnamed"

    # Limit length
    if len(safe_name) > 100:
        safe_name = safe_name[:97] + "..."

    return safe_name


def get_record_identifier(record: dict, fallback_index: int) -> str:
    """
    Get a suitable identifier for the record

    Args:
        record: JSON record
        fallback_index: Fallback index if no ID found

    Returns:
        String identifier for the record
    """
    # Try different ID fields in order of preference
    id_fields = ["id", "_id", "record_id", "uuid", "key"]

    for field in id_fields:
        if field in record and record[field]:
            return str(record[field])

    # Try to extract from common nested structures
    if "metadata" in record and isinstance(record["metadata"], dict):
        for field in id_fields:
            if field in record["metadata"] and record["metadata"][field]:
                return str(record["metadata"][field])

    # Use fallback
    return f"record_{fallback_index:06d}"


def convert_jsonl_to_files(
    input_file: str,
    output_dir: str,
    prefix: str = "",
    include_index: bool = True,
    indent: int = 2,
    limit: int = None,
    overwrite: bool = False,
) -> int:
    """
    Convert JSONL file to individual JSON files

    Args:
        input_file: Path to input JSONL file
        output_dir: Directory to save individual JSON files
        prefix: Prefix for output filenames
        include_index: Include sequential index in filename
        indent: JSON indentation (None for compact)
        limit: Maximum number of files to create
        overwrite: Whether to overwrite existing files

    Returns:
        Number of files created
    """
    input_path = Path(input_file)
    output_path = Path(output_dir)

    if not input_path.exists():
        raise FileNotFoundError(f"Input file not found: {input_file}")

    # Create output directory
    output_path.mkdir(parents=True, exist_ok=True)

    logging.info(f"Converting JSONL to individual JSON files")
    logging.info(f"Input: {input_file}")
    logging.info(f"Output directory: {output_dir}")
    logging.info(f"Prefix: {prefix or '(none)'}")
    logging.info(f"Include index: {include_index}")
    logging.info(f"JSON indent: {indent}")
    logging.info(f"Limit: {limit or 'No limit'}")

    created_count = 0
    skipped_count = 0
    error_count = 0

    with open(input_path, "r") as f:
        # Get total lines for progress bar
        total_lines = sum(1 for _ in f)
        f.seek(0)

        process_limit = min(total_lines, limit) if limit else total_lines

        with tqdm(total=process_limit, desc="Converting") as pbar:
            for line_num, line in enumerate(f, 1):
                if limit and line_num > limit:
                    break

                try:
                    # Parse JSON
                    record = json.loads(line.strip())

                    # Get record identifier
                    record_id = get_record_identifier(record, line_num)
                    safe_id = sanitize_filename(record_id)

                    # Build filename
                    filename_parts = []

                    if prefix:
                        filename_parts.append(prefix)

                    if include_index:
                        filename_parts.append(f"{line_num:03d}")

                    filename_parts.append(safe_id)

                    filename = "_".join(filename_parts) + ".json"
                    output_file_path = output_path / filename

                    # Check if file exists and handle overwrite
                    if output_file_path.exists() and not overwrite:
                        logging.debug(f"Skipping existing file: {filename}")
                        skipped_count += 1
                        pbar.update(1)
                        continue

                    # Write JSON file
                    with open(output_file_path, "w") as out_f:
                        json.dump(record, out_f, indent=indent, ensure_ascii=False)

                    created_count += 1
                    logging.debug(f"Created: {filename}")

                except json.JSONDecodeError as e:
                    logging.warning(f"Skipping malformed JSON at line {line_num}: {e}")
                    error_count += 1
                except Exception as e:
                    logging.error(f"Error processing line {line_num}: {e}")
                    error_count += 1

                pbar.update(1)

    logging.info(f"✅ Conversion complete!")
    logging.info(f"   Created: {created_count} files")
    if skipped_count > 0:
        logging.info(f"   Skipped: {skipped_count} files (already exist)")
    if error_count > 0:
        logging.warning(f"   Errors: {error_count} records")

    return created_count


def main():
    parser = argparse.ArgumentParser(
        description="Convert JSONL file to individual JSON files with proper formatting",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Basic conversion with sequential numbering
  python jsonl_to_individual_files.py --input records.jsonl --output individual_files/

  # With custom prefix and compact JSON
  python jsonl_to_individual_files.py --input data.jsonl --output results/ --prefix experiment_1 --indent 0

  # No index numbers, pretty formatting
  python jsonl_to_individual_files.py --input large_dataset.jsonl --output pretty_files/ --no-index --indent 4

  # Convert only first 50 records, overwrite existing
  python jsonl_to_individual_files.py --input huge.jsonl --output test/ --limit 50 --overwrite
        """,
    )

    parser.add_argument("--input", "-i", type=str, required=True, help="Input JSONL file")

    parser.add_argument(
        "--output", "-o", type=str, required=True, help="Output directory for individual JSON files"
    )

    parser.add_argument(
        "--prefix", type=str, default="", help="Prefix for output filenames (default: none)"
    )

    parser.add_argument(
        "--no-index", action="store_true", help="Do not include sequential index in filenames"
    )

    parser.add_argument(
        "--indent", type=int, default=2, help="JSON indentation (0 for compact, default: 2)"
    )

    parser.add_argument(
        "--limit", type=int, default=None, help="Maximum number of records to convert (for testing)"
    )

    parser.add_argument("--overwrite", action="store_true", help="Overwrite existing files")

    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose logging")

    args = parser.parse_args()

    setup_logging(args.verbose)

    try:
        created_count = convert_jsonl_to_files(
            input_file=args.input,
            output_dir=args.output,
            prefix=args.prefix,
            include_index=not args.no_index,
            indent=args.indent if args.indent > 0 else None,
            limit=args.limit,
            overwrite=args.overwrite,
        )

        print(f"\n✅ Successfully converted {created_count} records")
        print(f"📁 Files saved to: {args.output}")

        return 0

    except Exception as e:
        logging.error(f"Conversion failed: {e}")
        return 1


if __name__ == "__main__":
    exit(main())