"""
Data Mixing Service

Mixes fine-tune pool samples with randomly sampled old training data
following best practices for incremental YOLO fine-tuning.
"""

import logging
import random
import shutil
from pathlib import Path
from typing import Any
from uuid import UUID

logger = logging.getLogger(__name__)

# Mixing ratios: (max_new_samples, old_data_multiplier)
# Fewer new samples require higher old data ratio to prevent catastrophic forgetting
MIXING_RATIOS: list[tuple[int, int]] = [
    (10, 50),   # <= 10 new samples: 50x old
    (50, 20),   # <= 50 new samples: 20x old
    (200, 10),  # <= 200 new samples: 10x old
    (500, 5),   # <= 500 new samples: 5x old
]

# Default multiplier for 500+ samples
DEFAULT_MULTIPLIER = 5

# Maximum old samples to include (cap for performance)
MAX_OLD_SAMPLES = 3000

# Minimum pool size for fine-tuning
MIN_POOL_SIZE = 50


def get_mixing_ratio(new_sample_count: int) -> int:
    """Determine old data multiplier based on new sample count.

    Args:
        new_sample_count: Number of new samples in the fine-tune pool.

    Returns:
        Multiplier for old data sampling.
    """
    for threshold, multiplier in MIXING_RATIOS:
        if new_sample_count <= threshold:
            return multiplier
    return DEFAULT_MULTIPLIER


def build_mixed_dataset(
    pool_document_ids: list[UUID],
    base_dataset_path: Path,
    output_dir: Path,
    seed: int = 42,
) -> dict[str, Any]:
    """Build a mixed dataset for fine-tuning.

    Combines ALL fine-tune pool samples with randomly sampled old data
    from the base training dataset.

    Args:
        pool_document_ids: Document IDs from the fine-tune pool.
        base_dataset_path: Path to the base training dataset directory.
        output_dir: Output directory for the mixed dataset.
        seed: Random seed for reproducible sampling.

    Returns:
        Dictionary with dataset info (data_yaml path, counts).
    """
    new_count = len(pool_document_ids)
    multiplier = get_mixing_ratio(new_count)
    old_target = min(new_count * multiplier, MAX_OLD_SAMPLES)

    logger.info(
        "Building mixed dataset: %d new samples, %dx multiplier, "
        "targeting %d old samples",
        new_count, multiplier, old_target,
    )

    # Create output directory structure
    output_dir.mkdir(parents=True, exist_ok=True)
    for split in ("train", "val"):
        (output_dir / "images" / split).mkdir(parents=True, exist_ok=True)
        (output_dir / "labels" / split).mkdir(parents=True, exist_ok=True)

    # Collect old training images from base dataset
    old_train_images = _collect_images(base_dataset_path / "images" / "train")
    old_val_images = _collect_images(base_dataset_path / "images" / "val")

    # Randomly sample old data
    rng = random.Random(seed)
    all_old_images = old_train_images + old_val_images
    if len(all_old_images) > old_target:
        sampled_old = rng.sample(all_old_images, old_target)
    else:
        sampled_old = all_old_images

    # Split old samples: 80% train, 20% val
    rng.shuffle(sampled_old)
    old_train_count = int(len(sampled_old) * 0.8)
    old_train = sampled_old[:old_train_count]
    old_val = sampled_old[old_train_count:]

    # Copy old samples to mixed dataset
    old_copied = 0
    for split_name, images in [("train", old_train), ("val", old_val)]:
        for img_path in images:
            label_path = _image_to_label_path(img_path)
            dst_img = output_dir / "images" / split_name / img_path.name
            dst_label = output_dir / "labels" / split_name / label_path.name
            if img_path.exists():
                shutil.copy2(img_path, dst_img)
                old_copied += 1
            if label_path.exists():
                shutil.copy2(label_path, dst_label)

    # Copy new pool samples (from base dataset, identified by document_id prefix)
    # Pool documents go into train split (80%) and val split (20%)
    pool_id_strs = {str(doc_id) for doc_id in pool_document_ids}
    new_images = _find_pool_images(base_dataset_path, pool_id_strs)

    rng.shuffle(new_images)
    new_train_count = int(len(new_images) * 0.8)
    new_train = new_images[:new_train_count]
    new_val = new_images[new_train_count:]

    new_copied = 0
    for split_name, images in [("train", new_train), ("val", new_val)]:
        for img_path in images:
            label_path = _image_to_label_path(img_path)
            dst_img = output_dir / "images" / split_name / img_path.name
            dst_label = output_dir / "labels" / split_name / label_path.name
            if img_path.exists() and not dst_img.exists():
                shutil.copy2(img_path, dst_img)
                new_copied += 1
            if label_path.exists() and not dst_label.exists():
                shutil.copy2(label_path, dst_label)

    # Generate data.yaml
    from shared.fields import FIELD_CLASSES

    yaml_path = output_dir / "data.yaml"
    yaml_content = (
        f"path: {output_dir.absolute()}\n"
        f"train: images/train\n"
        f"val: images/val\n"
        f"\n"
        f"nc: {len(FIELD_CLASSES)}\n"
        f"names: {list(FIELD_CLASSES.values())}\n"
    )
    yaml_path.write_text(yaml_content)

    total_images = old_copied + new_copied
    logger.info(
        "Mixed dataset built: %d old + %d new = %d total images",
        old_copied, new_copied, total_images,
    )

    return {
        "data_yaml": str(yaml_path),
        "total_images": total_images,
        "old_images": old_copied,
        "new_images": new_copied,
        "mixing_ratio": multiplier,
    }


def _collect_images(images_dir: Path) -> list[Path]:
    """Collect all image files from a directory."""
    if not images_dir.exists():
        return []
    return sorted(images_dir.glob("*.png")) + sorted(images_dir.glob("*.jpg"))


def _image_to_label_path(image_path: Path) -> Path:
    """Convert image path to corresponding label path."""
    labels_dir = image_path.parent.parent.parent / "labels" / image_path.parent.name
    return labels_dir / image_path.with_suffix(".txt").name


def _find_pool_images(
    base_dataset_path: Path,
    pool_doc_ids: set[str],
) -> list[Path]:
    """Find images in base dataset that belong to pool documents."""
    images: list[Path] = []
    for split in ("train", "val", "test"):
        split_dir = base_dataset_path / "images" / split
        if not split_dir.exists():
            continue
        for img_path in split_dir.iterdir():
            if not img_path.is_file():
                continue
            # Image filenames are like: {doc_id}_page{N}.png
            doc_id = img_path.stem.rsplit("_page", 1)[0]
            if doc_id in pool_doc_ids:
                images.append(img_path)
    return images
