<?php
require_once __DIR__ . '/../lib/labdb.class.php';
require_once __DIR__ . '/../lib/job.class.php';
require_once __DIR__ . '/../lib/experiment.class.php';
require_once __DIR__ . '/../lib/parametercatalog.class.php';

final class LabWorker {
    private PDO $db;
    private array $config;
    private $lock;
    private string $activeLog = '';

    public function __construct(?array $config = null) {
        $this->config = $config ?? lab_config();
        $this->db = LabDb::connection();
    }

    public function acquire(): bool {
        $path = rtrim($this->config['jobs'], '/') . '/.worker.lock';
        if (!is_dir(dirname($path)) && !mkdir(dirname($path), 0750, true) && !is_dir(dirname($path))) return false;
        $this->lock = fopen($path, 'c');
        return $this->lock !== false && flock($this->lock, LOCK_EX | LOCK_NB);
    }

    public function recover(): void {
        $now = gmdate('c');
        $this->db->exec("UPDATE jobs SET status='interrupted',phase='interrupted',error=COALESCE(error,'Worker stopped before job completion'),finished_at=" . $this->db->quote($now) . ",updated_at=" . $this->db->quote($now) . ',pid=NULL WHERE status=\'running\'');
    }

    public function claim(): ?array {
        $path = $this->config['runtime_lock'];
        $runtime = fopen($path, is_file($path) ? 'r' : 'c');
        if ($runtime === false) return null;
        if (!flock($runtime, LOCK_SH | LOCK_NB)) { fclose($runtime); return null; }
        try {
            $this->db->exec('BEGIN IMMEDIATE');
            if ($this->db->query("SELECT 1 FROM jobs WHERE status='running' LIMIT 1")->fetchColumn()) { $this->db->commit(); return null; }
            $row = $this->db->query("SELECT * FROM jobs WHERE status='queued' ORDER BY created_at LIMIT 1")->fetch();
            if (!$row) { $this->db->commit(); return null; }
            $now = gmdate('c');
            $s = $this->db->prepare("UPDATE jobs SET status='running',phase='starting',started_at=?,updated_at=? WHERE id=? AND status='queued'");
            $s->execute([$now, $now, $row['id']]);
            if (!$s->rowCount()) { $this->db->commit(); return null; }
            $this->db->commit();
            try { $row['spec'] = json_decode($row['spec'], true, 512, JSON_THROW_ON_ERROR); }
            catch (JsonException $e) { $row['spec_error'] = 'Stored job spec is not valid JSON.'; $row['spec'] = null; }
            return $row;
        } catch (Throwable $e) { if ($this->db->inTransaction()) $this->db->rollBack(); throw $e; }
        finally { fclose($runtime); }
    }

    public function run(array $job): void {
        if (!empty($job['spec_error'])) { $this->finish($job['id'],'failed',$job['spec_error']); return; }
        if (!in_array($job['kind'] ?? '', ['model','planet'],true)) { $this->finish($job['id'],'failed','Unsupported job kind.'); return; }
        $validation = ParameterCatalog::validate($job['spec']);
        if (empty($validation['ok'])) { $this->finish($job['id'], 'failed', 'Stored job spec failed validation: ' . json_encode($validation['errors'] ?? [])); return; }
        $job['spec'] = $validation['spec'];
        $dir = Job::directory(['id' => $job['id']]);
        if (!is_dir($dir) && !mkdir($dir, 0750, true) && !is_dir($dir)) { $this->finish($job['id'], 'failed', 'Could not create the job directory.'); return; }
        $jobFile = $dir . '/job.json';
        $captured=Job::validatedMapCaptures($job);
        if(empty($captured['ok'])){$this->finish($job['id'],'failed',$captured['error']);return;}
        $mapSnapshot=Maps::writeSnapshots($captured['data'],$dir);
        if(empty($mapSnapshot['ok'])){$this->finish($job['id'],'failed','Map snapshot is unavailable: '.json_encode($mapSnapshot['errors']??[]));return;}
        $scientificSpec = $job['spec'];
        $scientificSpec['parameters'] = (object)$scientificSpec['parameters'];
        $this->atomicJson($jobFile, ['id'=>$job['id'], 'name'=>$job['name'], 'kind'=>$job['kind'], 'spec'=>$scientificSpec, 'directory'=>$dir, 'map_snapshots'=>$mapSnapshot['data']]);
        $configFile = getenv('EXOPLASIM_CONFIG') ?: $this->config['project_root'] . '/config/web.json';
        $configFile = realpath($configFile);
        if ($configFile === false) { $this->finish($job['id'],'failed','Configuration file does not exist.'); return; }
        $scriptRoot = $this->config['science_scripts'] ?? ($this->config['project_root'] . '/scripts/web');
        $script = $scriptRoot . ($job['kind'] === 'planet' ? '/planet.py' : '/model.py');
        $command = ['setsid', $this->config['python'], $script, '--config', $configFile, '--job', $jobFile];
        if (!is_file($script)) { $this->finish($job['id'],'failed','Science task script does not exist.'); return; }
        if (!is_file($this->config['python']) || !is_executable($this->config['python'])) { $this->finish($job['id'],'failed','Configured Python executable is unavailable.'); return; }
        $logPath = $dir . '/worker.log'; $this->activeLog = $logPath;
        $log = fopen($logPath, 'ab');
        if (!$log) { $this->finish($job['id'], 'failed', 'Could not open the job log.'); return; }
        if (disk_free_space($dir) < (int)$this->config['minimum_free_bytes']) { fclose($log); $this->finish($job['id'], 'failed', 'Insufficient free disk space.'); return; }
        $pipes = [];
        $environment=getenv(); if(!is_array($environment))$environment=[]; $environment['EXOPLASIM_CONFIG']=$configFile;
        $proc = proc_open($command, [0=>['pipe','r'], 1=>['pipe','w'], 2=>['file',$logPath,'a']], $pipes, $this->config['project_root'], $environment);
        fclose($log);
        if (!is_resource($proc)) { $this->finish($job['id'], 'failed', 'Could not start the science task.'); return; }
        fclose($pipes[0]); stream_set_blocking($pipes[1], false);
        $status = proc_get_status($proc); $pid = (int)$status['pid'];
        $this->update($job['id'], ['pid'=>$pid,'phase'=>'running']);
        $started = microtime(true); $buffer = ''; $lastCycle = 0; $datasetCount = 0; $error = null; $cancelled = false; $timedOut = false; $terminationAt = null; $exitCode = null;
        while (true) {
            $status = proc_get_status($proc);
            $chunk = stream_get_contents($pipes[1]);
            if ($chunk !== false && $chunk !== '') {
                $buffer .= $chunk;
                try { $this->events($job, $buffer, $lastCycle, $datasetCount, $error); }
                catch (Throwable $e) { $error=$e->getMessage(); if($terminationAt===null){$terminationAt=microtime(true);$this->signalGroup($pid,SIGTERM);} }
            }
            $current = $this->db->prepare('SELECT cancel_requested FROM jobs WHERE id=?'); $current->execute([$job['id']]);
            if ($error !== null && $terminationAt === null) { $terminationAt=microtime(true); $this->signalGroup($pid,SIGTERM); }
            if ((int)$current->fetchColumn() && !$cancelled) { $cancelled = true; $terminationAt = microtime(true); $this->signalGroup($pid, SIGTERM); }
            if (microtime(true) - $started > (int)$this->config['job_timeout_seconds'] && !$timedOut) { $timedOut = true; $terminationAt = microtime(true); $this->signalGroup($pid, SIGTERM); }
            if (disk_free_space($dir) < (int)$this->config['minimum_free_bytes'] && $error === null) { $error = 'Disk space fell below the configured floor.'; $terminationAt = microtime(true); $this->signalGroup($pid, SIGTERM); }
            if (!$status['running']) { $exitCode = $status['exitcode']; break; }
            if ($terminationAt !== null && microtime(true)-$terminationAt > 5) $this->signalGroup($pid, SIGKILL);
            usleep(200000);
        }
        $tail = stream_get_contents($pipes[1]); if ($tail !== false) { $buffer .= $tail; try { $this->events($job, $buffer, $lastCycle, $datasetCount, $error); } catch (Throwable $e) { $error=$e->getMessage(); } }
        if ($buffer !== '') file_put_contents($logPath,$buffer,FILE_APPEND);
        fclose($pipes[1]); $closedExit = proc_close($proc); $exit = $exitCode !== null && $exitCode >= 0 ? $exitCode : $closedExit;
        if ($cancelled) $this->finish($job['id'],'cancelled',null,$lastCycle);
        elseif ($timedOut) $this->finish($job['id'],'failed','Job exceeded its time limit.',$lastCycle);
        elseif ($error !== null) $this->finish($job['id'],'failed',$error,$lastCycle);
        elseif ($exit !== 0) $this->finish($job['id'],'failed',$this->tail($logPath),$lastCycle);
        elseif ($lastCycle < ($job['kind']==='planet'?1:(int)$job['total_cycles']) || $datasetCount < ($job['kind']==='planet'?1:(int)$job['total_cycles'])) $this->finish($job['id'],'failed','Science task exited without publishing all expected cycle results.',$lastCycle);
        else $this->finish($job['id'],'completed',null,$job['kind']==='planet'?1:(int)$job['total_cycles']);
    }

    private function events(array $job, string &$buffer, int &$lastCycle, int &$datasetCount, ?string &$error): void {
        while (($at = strpos($buffer, "\n")) !== false) {
            $line = substr($buffer, 0, $at); $buffer = substr($buffer, $at + 1);
            file_put_contents($this->activeLog,$line."\n",FILE_APPEND);
            $event = json_decode($line, true); if (!is_array($event)) continue;
            if (isset($event['error'])) $error = (string)$event['error'];
            if (isset($event['cycle'])) {
                $lastCycle = max($lastCycle, (int)$event['cycle']);
                $total = max(1, (int)($event['cycles'] ?? $job['total_cycles']));
                $this->update($job['id'], ['completed_cycles'=>$lastCycle,'progress'=>100.0*$lastCycle/$total,'phase'=>isset($event['dataset'])?'postprocessing':'running']);
                if (isset($event['dataset'])) {
                    if (!empty($event['planet'])) $this->publishPlanet($job,$event);
                    else $this->publishIndex($job, $event);
                    $datasetCount++;
                }
            }
            if (($event['planet'] ?? false) && isset($event['spec'])) $this->saveGeneratedExperiment($job, $event['spec']);
        }
    }

    private function publishPlanet(array $job,array $event): void {
        if(empty($event['dataset'])||!preg_match('/^[a-f0-9]{32}$/',$event['dataset'])) throw new RuntimeException('Planet task returned an invalid dataset ID.');
        $base=Job::directory(['id'=>$job['id']]).'/artifacts';$path=$base.'/datasets/'.$event['dataset'].'/index.json';
        $index=json_decode((string)@file_get_contents($path),true);
        if(!is_array($index)||($index['id']??null)!==$event['dataset']) throw new RuntimeException('Planet dataset index is missing or does not match its ID.');
        $this->atomicJson($base.'/index.json',['datasets'=>[['id'=>$event['dataset'],'label'=>'Generated planet surface']],'latest'=>$event['dataset']]);
    }

    private function saveGeneratedExperiment(array $job, array $spec): void {
        $valid = ParameterCatalog::validate($spec);
        if (empty($valid['ok'])) throw new RuntimeException('Generated experiment failed validation: ' . json_encode($valid['errors'] ?? []));
        $saved = Experiment::save($job['owner'], ['name'=>$job['name'],'spec'=>$valid['spec']]);
        if (empty($saved['ok'])) throw new RuntimeException('Could not save generated experiment.');
        $s=$this->db->prepare('UPDATE jobs SET experiment_id=? WHERE id=?'); $s->execute([$saved['data']['id'],$job['id']]);
    }

    private function publishIndex(array $job, array $event): void {
        if (empty($event['dataset']) || !preg_match('/^[a-f0-9]{32}$/',$event['dataset'])) return;
        $root=Job::directory(['id'=>$job['id']]) . '/artifacts';
        $datasetIndex=$root.'/datasets/'.$event['dataset'].'/index.json';
        $dataset=json_decode((string)@file_get_contents($datasetIndex),true);
        if (!is_array($dataset)||($dataset['id']??null)!==$event['dataset']) throw new RuntimeException('Science task dataset index is missing or invalid.');
        $path=$root.'/index.json';
        $index=is_file($path)?json_decode((string)file_get_contents($path),true):['datasets'=>[],'latest'=>null];
        $index['datasets'][]=['id'=>$event['dataset'],'label'=>'Cycle '.(int)($event['cycle']??1)]; $index['latest']=$event['dataset'];
        $this->atomicJson($path,$index);
    }

    private function signalGroup(int $pid, int $signal): void { if (function_exists('posix_kill')) @posix_kill(-$pid,$signal); }
    private function update(string $id, array $values): void { $sets=[]; $args=[]; foreach($values as $key=>$value){$sets[]=$key.'=?';$args[]=$value;} $sets[]='updated_at=?';$args[]=gmdate('c');$args[]=$id;$s=$this->db->prepare('UPDATE jobs SET '.implode(',',$sets).' WHERE id=?');$s->execute($args); }
    private function finish(string $id,string $status,?string $error,?int $cycles=null): void { $s=$this->db->prepare('UPDATE jobs SET status=?,phase=?,progress=CASE WHEN ?=\'completed\' THEN 100 ELSE progress END,completed_cycles=COALESCE(?,completed_cycles),error=?,pid=NULL,finished_at=?,updated_at=? WHERE id=?');$s->execute([$status,$status,$status,$cycles,$error,gmdate('c'),gmdate('c'),$id]); }
    private function tail(string $path): string { $file=@fopen($path,'rb');if(!$file)return 'Science task failed.';$size=fstat($file)['size']??0;fseek($file,max(0,$size-8192));$data=fread($file,8192);fclose($file);return $data===false?'Science task failed.':$data; }
    private function atomicJson(string $path,array $data): void { $tmp=$path.'.'.bin2hex(random_bytes(4)).'.tmp';file_put_contents($tmp,json_encode($data,JSON_THROW_ON_ERROR));rename($tmp,$path); }

    public function loop(): void {
        if (!$this->acquire()) return;
        $this->recover();
        while (true) { $job=$this->claim(); if (!$job) { usleep((int)($this->config['poll_seconds']??2)*1000000); continue; } try { $this->run($job); } catch (Throwable $e) { $this->finish($job['id'],'failed',$e->getMessage()); } }
    }
}

if (PHP_SAPI === 'cli' && realpath($_SERVER['SCRIPT_FILENAME'] ?? '') === __FILE__) (new LabWorker())->loop();