"""Preparation only: Docker save uses daemon-root temporary space, not just output space."""
import contextlib
import json
import os
from pathlib import Path
import shutil
import signal
import subprocess
import time

from ..core import absolute, docker, lock, read_json, require, sha256, write_json

GIB=1024**3
TEST_FLOOR = 10 * GIB
TEST_STOP = 12 * GIB
OUTPUT_RESERVE = GIB


def required_space(images, parent=None):
    sizes=sum(x['Size'] for x in images)
    if len(images)==2 and parent and all(x.get('Parent')==parent['Id'] for x in images):
        sizes-=parent['Size']
    return sizes+GIB  # tar headers and conservative workspace margin


def plan(references, output, *, test_preparation=False):
    output=absolute(output,no_symlinks=True)
    images=[json.loads(docker('image','inspect',r)[1])[0] for r in references]
    parent=None
    if len(images)==2 and images[0].get('Parent') and images[0].get('Parent')==images[1].get('Parent'):
        parent=json.loads(docker('image','inspect',images[0]['Parent'])[1])[0]
    _,root=docker('info','--format','{{.DockerRootDir}}');root=absolute(root.strip(),no_symlinks=True)
    ancestor=output.parent
    while not ancestor.exists():ancestor=ancestor.parent
    amount=required_space(images,parent)
    same=root.stat().st_dev==ancestor.stat().st_dev
    # This opt-in belongs to test image preparation, not installation/rollback or other sites.
    floor = TEST_FLOOR if test_preparation else 15 * GIB
    stop = TEST_STOP if test_preparation else 16 * GIB
    needed=amount*(2 if same else 1)+(stop if test_preparation else floor)
    result=dict(images=[x['Id'] for x in images],docker_root=str(root),
                docker_free=shutil.disk_usage(root).free,docker_required=needed,
                output_free=shutil.disk_usage(ancestor).free,output_required=amount+OUTPUT_RESERVE,
                reserve=floor,stop_reserve=stop,temp_budget=amount,
                policy='test-preparation-10-12' if test_preparation else 'legacy-preparation',
                same_filesystem=same)
    result['ready']=result['docker_free']>=needed and result['output_free']>=result['output_required']
    return result


def export(references, output, *, test_preparation=False, lock_directory=None, guard=None):
    output=absolute(output,no_symlinks=True)
    with lock(lock_directory or output.parent):
        return _export(references, output, test_preparation=test_preparation, guard=guard)


def _export(references, output, *, test_preparation, guard):
    p=plan(references,output,test_preparation=test_preparation)
    result_path = output.with_suffix('.result.json')
    if output.exists():
        require(result_path.is_file(),'OUTPUT_EXISTS','불완전 자료는 보존합니다. 재시도는 새 출력 경로로 명시하세요.')
        result = read_json(result_path)
        require(result['images'] == p['images'] and output.stat().st_size == result['bytes']
                and sha256(output) == result['sha256'], 'EXPORT_RESULT_CHANGED', '완료 자료의 입력/지문이 달라 재사용하지 않습니다.')
        return dict(result, reused=True)
    require(p['ready'],'DOCKER_EXPORT_SPACE',
            'Docker root 임시 예산과 선제 중단 여유, 출력 공간이 부족합니다. 보호 검사를 우회하지 않습니다.')
    output.parent.mkdir(parents=True,exist_ok=True)
    write_json(output.with_suffix('.intent.json'),p,new=True)
    if guard:
        guard()
    with output.with_suffix('.stderr.log').open('xb') as error, output.with_suffix('.samples.jsonl').open('x') as samples:
        os.chmod(error.name,0o600)
        os.chmod(samples.name,0o600)
        # Pin the inspected IDs; a tag moving after preflight must not change the export.
        proc=subprocess.Popen(['docker','--host','unix:///var/run/docker.sock','save','-o',str(output),*p['images']],
                              stdout=subprocess.DEVNULL,stderr=error,start_new_session=True)
        started = time.monotonic()
        deadline=started+1800
        previous = None
        minimum = p['docker_free']
        failure = None
        try:
            while proc.poll() is None:
                now = time.monotonic()
                free = shutil.disk_usage(p['docker_root']).free
                out_free = shutil.disk_usage(output.parent).free
                written = output.stat().st_size if output.exists() else 0
                minimum = min(minimum, free)
                samples.write(json.dumps(dict(elapsed=now-started, root_free=free,
                                              output_free=out_free, written=written))+'\n')
                samples.flush()
                require(free>=p['stop_reserve'], 'DOCKER_EXPORT_RESERVE',
                        '최저 보호선 전에 이번 반출 client만 중단합니다. daemon은 변경하지 않습니다.')
                if previous and free < p['stop_reserve'] + 2*GIB:
                    rate = max(0, previous[1]-free) / max(now-previous[0], .1)
                    require(free-rate*2>=p['reserve'], 'DOCKER_EXPORT_RAPID_DRAIN',
                            '급격한 공간 감소로 보호선 침범이 예상돼 반출만 중단합니다.')
                previous = (now, free)
                require(out_free>=max(p['temp_budget']-written,0)+OUTPUT_RESERVE,
                        'DOCKER_EXPORT_OUTPUT_SPACE','출력 디스크의 남은 쓰기 예산이 부족합니다.')
                require(now<deadline,'DOCKER_EXPORT_TIMEOUT','반출 제한시간을 초과했습니다.')
                if guard:
                    guard()
                time.sleep(1)
            require(proc.returncode==0,'DOCKER_EXPORT_FAILED','반출 실패. 부분 자료는 보존합니다.')
        except BaseException as exc:
            failure = getattr(exc, 'code', type(exc).__name__)
            raise
        finally:
            if proc.poll() is None:
                with contextlib.suppress(ProcessLookupError):
                    os.killpg(proc.pid,signal.SIGTERM)
                try:proc.wait(timeout=15)
                except subprocess.TimeoutExpired:
                    with contextlib.suppress(ProcessLookupError):
                        os.killpg(proc.pid,signal.SIGKILL)
                    proc.wait(timeout=5)
            if failure:
                write_json(output.with_suffix('.failure.json'),dict(code=failure, exit_code=proc.returncode,
                           min_root_free=minimum, partial_preserved=output.exists()),new=True)
    result=dict(images=p['images'],bytes=output.stat().st_size,sha256=sha256(output),
                min_root_free=minimum, elapsed=time.monotonic()-started, policy=p['policy'])
    write_json(result_path,result,new=True)
    return result


def export_many(jobs, *, test_preparation=False, lock_directory=None, guard=None):
    """One blocking save at a time; a failure stops the sequence, completed files are reusable."""
    outputs = [absolute(output, no_symlinks=True) for _, output in jobs]
    require(len(set(outputs)) == len(outputs), 'EXPORT_DUPLICATE', '같은 출력 경로를 중복 지정하지 마세요.')
    return [export(refs, output, test_preparation=test_preparation,
                   lock_directory=lock_directory, guard=guard) for refs, output in jobs]