"""Run `python -m app.worker`; add replicas to scale parsing independently of API."""
import json
import logging
import os
from pathlib import Path
import signal
import subprocess
import sys
import tempfile
import time
from sqlalchemy import update
from .config import settings
from .cleanup import cleanup
from .db import Job, claim, initialize, sessions
from .storage import storage

log = logging.getLogger(__name__)
stopping = False

def process(job):
    config = settings()
    prefix = f"{job['id']}/attempt-{job['attempt']}"
    try:
        with tempfile.TemporaryDirectory(prefix='schdoc-') as directory:
            root = Path(directory)
            source = root / 'input.SchDoc'
            comparison = job['input_key'].endswith('/comparison.json')
            if comparison:
                manifest = json.loads(storage().get(job['input_key']))
                source = root / 'comparison.json'
                source.write_text(json.dumps(manifest), encoding='utf-8')
                for side in ('old','new'):
                    (root / f'{side}.SchDoc').write_bytes(storage().get(f"{job['id']}/{side}.SchDoc"))
            else:
                source.write_bytes(storage().get(job['input_key']))
            env = {k: v for k, v in os.environ.items() if k in ('PATH', 'LANG', 'LC_ALL', 'PYTHONPATH')}
            # Drain no untrusted output into logs. Kill/wait is handled by subprocess.run.
            subprocess.run([sys.executable, '-m', 'app.compare_task' if comparison else 'app.parse_task', str(source), str(root), job['filename'], job['id']],
                           timeout=config.parse_timeout, check=True, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            result = json.loads((root / 'result.json').read_text())
            if result['ok']:
                storage().put(prefix + '/data.json', (root / 'data.json').read_bytes(), 'application/json')
                for name in (('old.svg','new.svg') if comparison else ('drawing.svg',)):
                    storage().put(prefix + '/' + name, (root / name).read_bytes(), 'image/svg+xml')
                values = dict(status='completed', output_prefix=prefix, summary=json.dumps(result['summary']))
            else:
                values = dict(status='failed', error_code=result['code'], error_message=result['message'])
    except subprocess.TimeoutExpired:
        values = dict(status='failed', error_code='parse_timeout', error_message='Schematic exceeded the parsing time limit')
    except subprocess.CalledProcessError:
        values = dict(status='failed', error_code='parse_resource_limit', error_message='Parser exited before completion (resource limit or invalid input)')
    except Exception:
        log.exception('Worker infrastructure failure for %s', job['id'])
        values = dict(status='failed', error_code='processing_error', error_message='Storage or processing service unavailable')
    with sessions().begin() as db:
        # An expired worker may finish late. Its outputs must never replace a newer attempt.
        db.execute(update(Job).where(Job.id == job['id'], Job.attempt == job['attempt'], Job.status == 'processing')
                   .values(**values, lease_until=None, updated_at=time.time()))

def main():
    global stopping
    logging.basicConfig(level=logging.INFO)
    def stop(*_):
        global stopping
        stopping = True
    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)
    initialize()
    last_cleanup=0
    while not stopping:
        if time.monotonic()-last_cleanup>60:
            cleanup(); last_cleanup=time.monotonic()
        job = claim()
        if job: process(job)
        else: time.sleep(0.5)

if __name__ == '__main__': main()