#!/usr/bin/env python3
"""
Sample rows from vocal stems JSONL file while preserving order.
Supports both ratio-based and count-based sampling.
"""

import json
import argparse
import math
import random
from pathlib import Path
from tqdm import tqdm


def sample_by_count(input_path: Path, output_path: Path, count: int, random_seed: int = 42):
    """
    Sample a specific number of rows uniformly distributed across the file.
    Preserves original order in output.
    """
    random.seed(random_seed)

    # First pass: count total lines
    print("Counting total lines...")
    with open(input_path, "r") as f:
        total_lines = sum(1 for _ in tqdm(f, desc="Counting lines", unit=" lines"))

    if count >= total_lines:
        print(f"Requested count {count} >= total lines {total_lines}, copying all lines")
        count = total_lines
        sample_indices = set(range(total_lines))
    else:
        # Calculate uniform sampling indices
        step = total_lines / count
        sample_indices = set()
        for i in range(count):
            index = int(i * step + random.uniform(0, step))
            index = min(index, total_lines - 1)  # Ensure we don't exceed bounds
            sample_indices.add(index)

        # If we didn't get enough unique indices due to rounding, add more
        while len(sample_indices) < count:
            remaining = set(range(total_lines)) - sample_indices
            if remaining:
                sample_indices.add(random.choice(list(remaining)))
            else:
                break

    print(f"Sampling {len(sample_indices)} lines from {total_lines} total lines")

    # Second pass: collect sampled lines in order
    sampled_lines = []
    with open(input_path, "r") as f:
        for line_num, line in enumerate(
            tqdm(f, desc="Sampling lines", unit=" lines", total=total_lines)
        ):
            if line_num in sample_indices:
                sampled_lines.append((line_num, line.strip()))

    # Sort by original line number to preserve order
    sampled_lines.sort(key=lambda x: x[0])

    # Write sampled lines to output
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        for _, line in tqdm(sampled_lines, desc="Writing output", unit=" lines"):
            f.write(line + "\n")

    return len(sampled_lines)


def sample_by_ratio(input_path: Path, output_path: Path, ratio: float, random_seed: int = 42):
    """
    Sample a percentage of rows uniformly distributed across the file.
    Preserves original order in output.
    """
    random.seed(random_seed)

    # First pass: count total lines
    print("Counting total lines...")
    with open(input_path, "r") as f:
        total_lines = sum(1 for _ in tqdm(f, desc="Counting lines", unit=" lines"))

    target_count = int(total_lines * ratio)
    print(f"Target count: {target_count} ({ratio:.1%} of {total_lines})")

    if target_count <= 0:
        print("Error: Ratio too small, would result in 0 samples")
        return 0

    return sample_by_count(input_path, output_path, target_count, random_seed)


def sequential_sample(input_path: Path, output_path: Path, count: int):
    """
    Sample first N rows sequentially (for testing/debugging).
    """
    print(f"Taking first {count} rows sequentially")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    lines_written = 0

    with open(input_path, "r") as infile, open(output_path, "w") as outfile:
        for line_num, line in enumerate(tqdm(infile, desc="Sequential sampling", unit=" lines")):
            if line_num >= count:
                break
            outfile.write(line)
            lines_written += 1

    return lines_written


def main():
    parser = argparse.ArgumentParser(description="Sample rows from vocal stems JSONL file")
    parser.add_argument(
        "--input-path",
        type=str,
        default="/home/vibert/data/voice_designer/metas_v6_tr_vocal_stems.jsonl",
        help="Path to input vocal stems JSONL file",
    )
    parser.add_argument(
        "--output-path", type=str, required=True, help="Path to output sampled JSONL file"
    )

    # Sampling method (mutually exclusive)
    sampling_group = parser.add_mutually_exclusive_group(required=True)
    sampling_group.add_argument("--count", type=int, help="Sample specific number of rows")
    sampling_group.add_argument(
        "--ratio", type=float, help="Sample ratio (0.0 to 1.0, e.g., 0.1 for 10%%)"
    )
    sampling_group.add_argument(
        "--sequential", type=int, help="Take first N rows sequentially (for testing)"
    )

    parser.add_argument("--random-seed", type=int, default=42, help="Random seed for reproducibility")
    parser.add_argument("--quiet", action="store_true", help="Suppress verbose output")

    args = parser.parse_args()

    input_path = Path(args.input_path)
    output_path = Path(args.output_path)

    if not input_path.exists():
        print(f"Error: Input file {input_path} does not exist")
        return 1

    if not args.quiet:
        print(f"Input:  {input_path}")
        print(f"Output: {output_path}")
        print()

    try:
        if args.count:
            sampled_count = sample_by_count(input_path, output_path, args.count, args.random_seed)
        elif args.ratio:
            if not (0.0 < args.ratio <= 1.0):
                print("Error: Ratio must be between 0.0 and 1.0")
                return 1
            sampled_count = sample_by_ratio(input_path, output_path, args.ratio, args.random_seed)
        elif args.sequential:
            sampled_count = sequential_sample(input_path, output_path, args.sequential)

        print(f"\n{'='*50}")
        print(f"SAMPLING COMPLETE")
        print(f"{'='*50}")
        print(f"Sampled records: {sampled_count:,}")
        print(f"Output saved to: {output_path}")

        return 0

    except Exception as e:
        print(f"Error during sampling: {e}")
        return 1


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