"""Burst/HDR group detection for photo metadata cleanup."""

from __future__ import annotations

from collections import defaultdict
from datetime import datetime, timedelta
from typing import TYPE_CHECKING

from tools.filenames import naive_wall_time

if TYPE_CHECKING:
    from photo_metadata import MediaRecord


def photo_sort_key(record: MediaRecord) -> tuple:
    assert record.adjusted_time is not None
    return (
        naive_wall_time(record.adjusted_time),
        record.subsec if record.subsec is not None else -1,
        record.sequence if record.sequence is not None else -1,
        record.path.name.lower(),
    )


def group_sort_key(record: MediaRecord) -> tuple:
    timestamp = naive_wall_time(record.adjusted_time) if record.adjusted_time else datetime.max
    return (
        timestamp,
        record.subsec if record.subsec is not None else -1,
        record.sequence if record.sequence is not None else -1,
        str(record.path).lower(),
    )


def detect_photo_groups(records: list[MediaRecord], min_size: int = 2) -> list[list[MediaRecord]]:
    by_camera: dict[str, list[MediaRecord]] = defaultdict(list)
    for record in records:
        if record.is_image and record.adjusted_time is not None:
            by_camera[record.camera_key].append(record)
    return [
        group
        for camera_records in by_camera.values()
        for group in detect_groups_for_camera(sorted(camera_records, key=photo_sort_key), min_size)
    ]


def detect_groups_for_camera(records: list[MediaRecord], min_size: int) -> list[list[MediaRecord]]:
    by_second: dict[datetime, list[MediaRecord]] = defaultdict(list)
    for record in records:
        assert record.adjusted_time is not None
        by_second[naive_wall_time(record.adjusted_time).replace(microsecond=0)].append(record)

    seconds = sorted(by_second)
    spans: list[list[datetime]] = []
    for second in seconds:
        if not spans or second != spans[-1][-1] + timedelta(seconds=1):
            spans.append([])
        spans[-1].append(second)

    groups: list[list[MediaRecord]] = []
    for span in spans:
        if max(len(by_second[second]) for second in span) < 2:
            continue
        span_records = [record for second in span for record in by_second[second]]
        if len(span_records) >= 2 and all(record.sequence is not None for record in span_records):
            groups.extend(sequence_groups(span_records, min_size))
        else:
            groups.extend(timestamp_groups(span, by_second, min_size))
    return [group for group in groups if len(group) >= min_size]


def sequence_groups(records: list[MediaRecord], min_size: int) -> list[list[MediaRecord]]:
    groups: list[list[MediaRecord]] = []
    for record in sorted(records, key=photo_sort_key):
        if not groups or record.sequence <= groups[-1][-1].sequence:
            groups.append([])
        groups[-1].append(record)
    return [group for group in groups if len(group) >= min_size]


def timestamp_groups(
    span: list[datetime], by_second: dict[datetime, list[MediaRecord]], min_size: int
) -> list[list[MediaRecord]]:
    multi_indices = [index for index, second in enumerate(span) if len(by_second[second]) >= 2]
    if not multi_indices:
        return []

    clusters: list[list[int]] = [[multi_indices[0]]]
    for index in multi_indices[1:]:
        if index - clusters[-1][-1] > 2:
            clusters.append([])
        clusters[-1].append(index)

    groups: list[list[MediaRecord]] = []
    used_seconds: set[datetime] = set()
    for cluster in clusters:
        start, end = cluster[0], cluster[-1]
        if start and len(by_second[span[start - 1]]) == 1:
            start -= 1
        if end + 1 < len(span) and len(by_second[span[end + 1]]) == 1:
            end += 1
        seconds = [second for second in span[start : end + 1] if second not in used_seconds]
        used_seconds.update(seconds)
        records = [record for second in seconds for record in sorted(by_second[second], key=photo_sort_key)]
        if len(records) >= min_size:
            groups.append(records)
    return groups


def infer_group_subseconds(groups: list[list[MediaRecord]]) -> None:
    for group in groups:
        if any(record.subsec is not None for record in group):
            continue
        by_second: dict[datetime, list[MediaRecord]] = defaultdict(list)
        for record in group:
            if record.is_image and record.adjusted_time is not None:
                by_second[naive_wall_time(record.adjusted_time).replace(microsecond=0)].append(record)

        seconds = sorted(by_second)
        counts = {second: len(records) for second, records in by_second.items()}
        for index, second in enumerate(seconds):
            records = sorted(by_second[second], key=photo_sort_key)
            count = len(records)
            if count <= 1:
                continue
            if len(seconds) == 2:
                fps = max(counts.values())
            elif index == 0:
                fps = max(count, counts[seconds[1]])
            elif index == len(seconds) - 1:
                fps = max(count, counts[seconds[-2]])
            else:
                fps = count
            start_slot = fps - count if index == 0 and fps > count else 0
            for item_index, record in enumerate(records):
                record.inferred_subsec = int((start_slot + item_index) * 1000 / fps)