import argparse
import subprocess
import os
import sys
import datetime
import shutil as _shutil
import calendar

import numpy as np
from sortedcontainers import SortedDict
import tqdm
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from suvi import paths, vfs


def find_ffmpeg():
    """Locate ffmpeg: $FFMPEG first, then PATH, then the bundled Windows binary."""
    override = os.environ.get("FFMPEG")
    if override and os.path.exists(override):
        return override
    found = _shutil.which("ffmpeg")
    if found:
        return found
    bundled = os.path.join(os.path.dirname(paths.data_root()), "ffmpeg.exe")
    if os.path.exists(bundled):
        return bundled
    raise SystemExit("ffmpeg not found. Install it, or set $FFMPEG to its path.")


def encode_videos(args):
    """Encode composite JPGs into per-satellite videos.

    Gaps are padded with black frames and smoothed by ffmpeg's minterpolate, except
    where the gap exceeds `max_frame_interp`, which is skipped entirely.  That
    behaviour is what the bench's "today" variant reproduces.
    """
    years = args.years
    sources = args.sources
    ffmpeg_path = find_ffmpeg()
    project_root = args.out or os.path.dirname(paths.data_root())
    os.makedirs(project_root, exist_ok=True)

    for year in years:
        for source in sources:
            path_to_images = args.images or os.path.join(
            os.path.dirname(paths.data_root()), "composite", source)
            output_file = os.path.join(project_root, f"{source}_{year}_{args.suffix}.mp4")
            interp_file = os.path.join(project_root, f"{source}_{year}_interp_{args.suffix}.mp4")

            starttime = calendar.timegm(datetime.datetime(year, 1, 1, tzinfo=datetime.timezone.utc).timetuple())
            stoptime = calendar.timegm(datetime.datetime(year+1, 1, 1, tzinfo=datetime.timezone.utc).timetuple())
            min_file_size = 350000 # Detect and remove corrupted files by filtering by file size
            max_file_size = 450000
            encoding_crf = 16
            max_frame_interp = 120

            # Set up ffmpeg to stream images from an input pipe
            command_line = f'{ffmpeg_path} -y -f image2pipe -framerate 60 -i - -c:v libx264 -crf {encoding_crf} -preset veryfast {output_file}'
            print(command_line)
            p = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE, stdin=subprocess.PIPE)

            black_image = Image.fromarray(np.zeros((1080, 1920, 3), dtype='uint8'))

            files_by_time = SortedDict()
            for root, dirs, files in tqdm.tqdm(os.walk(path_to_images), desc="Finding and filtering files."):
                for f in files:
                    if f.endswith('.jpg'):
                        fpath = os.path.join(root, f)
                        fsize = os.path.getsize(fpath)
                        time_chunk = f.replace('-', '.').replace('_', '.').split('.')[1]
                        ftime = int(time_chunk)
                        if (fsize > min_file_size) and (fsize < max_file_size) and (ftime > starttime) and (ftime < stoptime):
                            files_by_time[ftime] = fpath

            difftimes = np.diff(files_by_time.keys())
            unique, counts = np.unique(difftimes, return_counts=True)
            interval = unique[0]
            assert np.sum((unique % interval) > 0) == 0 # Ensure all our timestamps align perfectly with our interval

            prevtime = files_by_time.peekitem(0)[0] - interval
            reliever = vfs.Reliever(label=f"encode {source}")
            for t, f in tqdm.tqdm(files_by_time.items(), desc="FFMPEG encoding images", total = len(files_by_time)):
                reliever.tick()
                framejump = (t - prevtime) // interval
                if framejump < max_frame_interp:
                    for i in range(framejump - 1): # Fill black frames to fill gaps, these will later be interpolated with ffmpeg
                        black_image.save(p.stdin, 'jpeg', quality = 95)
                else: # Skip past intervals that are too large to fill reasonably
                    print(f"Detected a frame gap of: {framejump}!")
                with open(f, 'rb') as fh:
                    p.stdin.write(fh.read())
                prevtime = t

            p.stdin.close() # Close the ffmpeg input pipe
            p.wait() # Wait for ffmpeg to finish encoding

            command_line = f'{ffmpeg_path} -y -i {output_file} -vf blackframe=0,metadata=select:key=lavfi.blackframe.pblack:value=95:function=less,minterpolate=mi_mode=mci:mc_mode=aobmc:me_mode=bidir:vsbmc=1:me=fss -c:a copy -c:v libx264 -crf {encoding_crf} -preset veryfast {interp_file}'
            print(command_line)

            pipe = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE).stdout
            pipe.read()
            pipe.close()

def build_parser():
    parser = argparse.ArgumentParser(
        description="Encode SUVI composite images into videos.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--years", type=int, nargs="+", default=[2024])
    parser.add_argument("--sources", nargs="+", default=["goes16", "goes18"])
    parser.add_argument("--images", default=None,
                        help="directory of composites (default: <root>/composite/<source>)")
    parser.add_argument("--out", default=None, help="directory for the video files")
    parser.add_argument("--start", type=int, default=None, help="Unix time, inclusive")
    parser.add_argument("--stop", type=int, default=None, help="Unix time, exclusive")
    parser.add_argument("--suffix", default="nofilt", help="tag in the output filenames")
    return parser


if __name__ == "__main__":
    # Guarded: this module is imported for find_ffmpeg(), and without the guard that
    # import kicked off a full-year encode as a side effect.
    sys.exit(encode_videos(build_parser().parse_args()) or 0)