"""Worker drain and verified history restore for rolling updates.

Control files are operator-owned, read-only mounts. They cannot be supplied in a
generation prompt. Existing history and output reads are never gated by drain.
"""
import hashlib
import json
import os
from pathlib import Path
import secrets
import sys
import uuid


def maintenance_enabled(path):
    try:
        data=Path(path).read_bytes()
        if len(data)>4096:return True
        value=json.loads(data)
        if value.get('schema')!=1 or not isinstance(value.get('maintenance'),bool):return True
        return value['maintenance']
    except (OSError,ValueError,TypeError):return True


def make_middleware(path,canary_token):
    from aiohttp import web
    @web.middleware
    async def gate(request,handler):
        controlled=(request.method=='GET' and request.path in ['/queue','/api/queue']) or (
                    request.method=='POST' and request.path in ['/prompt','/api/prompt'])
        supplied=request.headers.get('X-H3-Canary-Token','')
        bypass=bool(canary_token and secrets.compare_digest(supplied,canary_token))
        if controlled and not bypass and maintenance_enabled(path):
            return web.json_response({'error':{'type':'h3_worker_draining',
                'message':'Worker maintenance: existing jobs/results are retained; new jobs use the other worker.'}},status=503,
                headers={'Retry-After':'10'})
        # Authentication is still handled by the original middleware/handler.
        # A canary token only bypasses drain, never normal API authorization.
        return await handler(request)
    return gate


def restore_history(queue,path,expected_sha256,expected_source):
    p=Path(path)
    if p.stat().st_size>256*1024**2:raise ValueError('History snapshot exceeds the audited size bound')
    raw=p.read_bytes()
    if not expected_sha256 or hashlib.sha256(raw).hexdigest()!=expected_sha256:
        raise ValueError('History snapshot SHA256 mismatch')
    snapshot=json.loads(raw)
    if snapshot.get('schema')!=1 or snapshot.get('queue_empty') is not True or snapshot.get('source_container_id')!=expected_source:
        raise ValueError('History snapshot identity or drain state mismatch')
    history=snapshot.get('history')
    if not isinstance(history,dict):raise ValueError('History snapshot must contain a dictionary')
    for key,value in history.items():
        uuid.UUID(key)
        if not isinstance(value,dict) or value.get('status',{}).get('status_str') not in ['success','error']:
            raise ValueError('A nonterminal history entry cannot be imported')
    with queue.mutex:
        if queue.queue or queue.currently_running:
            raise ValueError('Never restore a snapshot into a worker already executing jobs')
        overlap=set(history)&set(queue.history)
        if any(history[k]!=queue.history[k] for k in overlap):raise ValueError('History identifier collision')
        queue.history.update(history)
    return len(history)


def register_worker_state():
    server=sys.modules.get('server')
    instance=getattr(getattr(server,'PromptServer',None),'instance',None)
    if instance is None or getattr(instance,'_h3_worker_state_registered',False):return
    control=os.environ.get('H3_WORKER_CONTROL_FILE','')
    if control:
        token=os.environ.get('H3_CANARY_TOKEN','')
        token_file=os.environ.get('H3_CANARY_TOKEN_FILE','')
        if token_file:
            token=Path(token_file).read_text().strip()
            if len(token)<16:raise ValueError('Canary token file is missing a valid operator token')
        instance.app.middlewares.append(make_middleware(control,token))
    snapshot=os.environ.get('H3_HISTORY_SNAPSHOT','')
    if snapshot:
        async def restore_before_listen(app):
            count=restore_history(instance.prompt_queue,snapshot,os.environ.get('H3_HISTORY_SNAPSHOT_SHA256',''),
                                  os.environ.get('H3_HISTORY_SOURCE_CONTAINER_ID',''))
            print(f'H3 verified history restored: {count} terminal jobs; no prompt resubmission',flush=True)
        instance.app.on_startup.append(restore_before_listen)
    instance._h3_worker_state_registered=True