from dataclasses import dataclass
from fractions import Fraction
from pathlib import Path
import hashlib
import av
from .config import MAX_DURATION_S, MAX_DIMENSION, MAX_FRAME_COUNT

class InputError(ValueError):
    pass

@dataclass(frozen=True)
class Video:
    path: Path
    sha256: str
    rate: Fraction
    pts_us: tuple
    duration_us: int
    size: tuple
    cfr: bool

    @property
    def duration_s(self):
        return self.duration_us / 1000000

def probe(path):
    with av.open(str(path), options={'err_detect': 'explode'}) as container:
        if len(container.streams.video) != 1:
            raise InputError('Use a file containing one video stream.')
        stream = container.streams.video[0]
        stream.codec_context.thread_count = 1
        rate = stream.average_rate
        if not rate or rate < 4:
            raise InputError('The input needs at least 4 source frames per second.')
        size = (stream.codec_context.width, stream.codec_context.height)
        if min(size) <= 0 or max(size) > MAX_DIMENSION:
            raise InputError('Video dimensions must be at most 4096 pixels per side.')
        times, zero, last_length = ([], None, None)
        for frame in container.decode(stream):
            if frame.pts is None:
                raise InputError('Every source frame needs a presentation timestamp.')
            zero = frame.pts if zero is None else zero
            t = Fraction(frame.pts - zero) * stream.time_base * 1000000
            if times and t <= times[-1]:
                raise InputError('Source timestamps must be strictly increasing.')
            if t >= MAX_DURATION_S * 1000000 or len(times) >= MAX_FRAME_COUNT:
                raise InputError(f'Video exceeds {MAX_DURATION_S:g} seconds or {MAX_FRAME_COUNT} frames.')
            times.append(t)
            last_length = Fraction(frame.duration) * stream.time_base * 1000000 if frame.duration else None
        if len(times) < 2:
            raise InputError('The video contains too few source frames.')
        tail = last_length or Fraction(1000000, rate)
        duration = round(times[-1] + tail)
        if duration > MAX_DURATION_S * 1000000:
            raise InputError(f'Video exceeds {MAX_DURATION_S:g} seconds.')
        cfr = all((abs(t - Fraction(i * 1000000, rate)) <= 1 for i, t in enumerate(times)))
    with path.open('rb') as f:
        digest = hashlib.file_digest(f, 'sha256').hexdigest()
    return Video(path, digest, rate, tuple(times), duration, size, cfr)

def check_synchronized(videos):
    main = videos['main']
    if len(videos) == 3:
        for role, video in videos.items():
            if video.rate != main.rate or video.pts_us != main.pts_us or video.duration_us != main.duration_us:
                raise InputError('Three camera files must have identical relative frame timestamps, frame counts and duration.')
        if not main.cfr:
            raise InputError('The three-view profile requires constant-rate synchronized source videos.')

def decode_selected(video, wanted, canonical=False):
    wanted = set(wanted)
    result = {}
    if not wanted:
        return result
    if min(wanted) < 0 or max(wanted) >= len(video.pts_us):
        raise InputError('A requested source frame is outside the video.')
    with av.open(str(video.path), options={'err_detect': 'explode'}) as container:
        stream = container.streams.video[0]
        stream.codec_context.thread_count = 1
        zero = None
        for i, frame in enumerate(container.decode(stream)):
            zero = frame.pts if zero is None else zero
            if i in wanted:
                if Fraction(frame.pts - zero) * stream.time_base * 1000000 != video.pts_us[i]:
                    raise InputError('Source frame timestamps changed while decoding.')
                use_canonical = canonical and frame.format.name == 'yuv420p' and (frame.width % 16 == 0) and (frame.height % 2 == 0) and (int(frame.colorspace) in (2, 5, 6)) and (int(frame.color_range) in (0, 1))
                if use_canonical:
                    from canonical_rgb import canonical_rgb
                    result[i] = canonical_rgb(frame)
                else:
                    result[i] = frame.to_image().convert('RGB')
            if i >= max(wanted):
                break
    if set(result) != wanted:
        raise InputError('A requested source frame could not be decoded.')
    return result

def frames_for(videos, wanted, mosaic='YAM'):
    frames = {role: decode_selected(video, wanted, canonical=len(videos) == 3) for role, video in videos.items()}
    if len(videos) == 1:
        return frames['main']
    from composers import compose_top_dominant_mosaic, compose_frame
    if mosaic == 'YAM':
        return {i: compose_top_dominant_mosaic(frames['main'][i], frames['left_wrist'][i], frames['right_wrist'][i]) for i in wanted}
    if mosaic == 'Dex':
        return {i: compose_frame(dict(ego=frames['main'][i], left_wrist=frames['left_wrist'][i], right_wrist=frames['right_wrist'][i]), ('ego', 'left_wrist', 'right_wrist')) for i in wanted}
    raise InputError('Choose a supported camera layout.')