#!/usr/bin/perl
use strict;
use warnings;
use DBI;
use Digest::SHA qw(sha256_hex);
use JSON;
use File::Basename qw(dirname basename);
use File::Copy qw(move);
use String::ShellQuote qw/shell_quote/;

my $db_host = $ENV{'DB_HOST'} || 'mariadb';
my $db_name = $ENV{'DB_NAME'} || 'videodetect';
my $db_user = $ENV{'DB_USER'} || 'videodetect';
my $db_pass = $ENV{'DB_PASSWORD'} || 'videodetect';  
my $tasks = $ENV{'TASK_COUNT'} || 1;

$SIG{INT} = sub { die "Interrupted\n"; };

my $dbh = DBI->connect(
    "DBI:mysql:database=$db_name;host=$db_host;mysql_enable_utf8mb4=1",
    $db_user,
    $db_pass,
    { mysql_auto_reconnect => 1 }
);

my $sth_claimtask = $dbh->prepare("
UPDATE tasks
SET status='IN_PROGRESS', updated_at=NOW(), assign_key=?, assigned_at=NOW()
WHERE status='PENDING' AND task_type='TRANSCODE' AND assign_key IS NULL
LIMIT 1
");
my $sth_gettask = $dbh->prepare("SELECT * FROM tasks WHERE assign_key=? AND status='IN_PROGRESS' LIMIT 1");
my $sth_getvideo = $dbh->prepare("SELECT * from videos WHERE id=?");
my $sth_updatevideo = $dbh->prepare("UPDATE videos SET file_path=?, file_size=?, file_hash=?, updated_at=NOW() WHERE id=?");
my $sth_finishtask = $dbh->prepare("UPDATE tasks SET status='COMPLETED', updated_at=NOW(), assign_key=NULL, assigned_at=NULL, results=? WHERE id=?");

for(my $i=0; $i<$tasks; $i++) {
    my $assign_key = "worker_" . int(rand(10000000));
    print "Assigning task with key: $assign_key\n";
    my $count = $sth_claimtask->execute($assign_key);
    if($count != 1) {
        print "Failed to claim task\n";
        last;
    }

    $sth_gettask->execute($assign_key);
    my $task = $sth_gettask->fetchrow_hashref();
    if(!ref $task) {
        print "Failed to fetch task\n";
        last;
    }

    my $video_id = $task->{video_id};
    $sth_getvideo->execute($video_id);
    my $video = $sth_getvideo->fetchrow_hashref();
    if ($video) {
        print "Processing video: $video->{id}\n";
    }

    my $results = transcode_video($video);
    if ($results->{success}) {
        my $file_path = $results->{output};
        my $file_size = -s $file_path;
        my $file_hash = $results->{newhash};
        delete $results->{newhash};
        $sth_updatevideo->execute($file_path, $file_size, $file_hash, $video_id);
        $sth_finishtask->execute(JSON->new->encode($results), $task->{id});
    }
}

sub get_hash {
    my ($file_path) = @_;
    # --- Compute SHA-256 hash of the file ---
    return undef unless -f $file_path && -r $file_path;
    open(my $fh, '<:raw', $file_path) or do { warn "Cannot open $file_path: $!"; return undef; };
    my $hash = Digest::SHA->new(256)->addfile($fh)->hexdigest;    close($fh);
    return $hash;
}

sub transcode_video {
    my ($video, $retry) = @_;

    my $file = $video ? $video->{file_path} : undef;

    # --- validate ---
    unless (defined $file && -f $file) {
        warn "SKIP  '" . ($file // 'undef') . "' not found\n";
        return { success => 0, reason => 'not_found' };
    }

    # --- derive paths ---
    my $dir  = dirname($file);
    my $base = basename($file);
    my $name = $base;
    $name =~ s/\.[^.]*$//;           # strip last extension ("foo.mkv" -> "foo")
    $name = $base if $name eq '';    # guard against dot-only names

    my $output = "$dir/$name.mp4";   # equals $file when it's already .mp4
    my $tmp    = "/tmp/${name}~.mp4";

    # --- transcode ---
    print "WORK  $file  →  $output\n";

    my $cmd = 'ffmpeg -y -nostdin '
            . '-hwaccel cuda '
            . ($retry ? '' : '-hwaccel_output_format cuda ')
            . '-i ' . shell_quote($file) . ' '
            . '-c:v h264_nvenc -preset p6 -rc vbr -cq 23 -b:v 0 '
            . ($retry ? '-pix_fmt yuv420p ' : '')
            . '-c:a aac -b:a 128k '
            . '-movflags +faststart '
            . shell_quote($tmp);

    my $log       = `$cmd`;
    my $exit_code = $? >> 8;

    if ($exit_code == 0) {
        my $newhash = get_hash($tmp);
        # success — delete original, move transcoded file into place
        unless (unlink($file)) {
            warn "  WARN  could not delete original '$file': $!\n";
        }
        unless (move($tmp, $output)) {
            warn "  FAIL  could not move '$tmp' -> '$output': $!\n";
            return { success => 0, reason => "move_failed: $!" };
        }
        print "  OK    original removed, output: $output\n";
        return { success => 1, output => $output, original => $file, newhash => $newhash };
    }

    # failure — clean up temp file, keep original
    unlink($tmp) if -e $tmp;
    warn "  FAIL  keeping original: $file\n";
    warn $log if defined $log && length $log;

    if($retry) {
        return { success => 0, reason => 'transcode_failed', exit_code => $exit_code };
    } else {
        return transcode_video($video, 1);  # retry with $retry set to 1
    }
}