"""Install the official pinned Linux Blender build with working CPU denoising."""
import argparse
import hashlib
import http.client
import json
import os
from pathlib import Path
import platform
import shutil
import subprocess
import tarfile
import tempfile
import urllib.request
import urllib.error
from urllib.parse import urlsplit

VERSION = '5.1.2'
ARCHIVE = f'blender-{VERSION}-linux-x64.tar.xz'
URL = f'https://download.blender.org/release/Blender5.1/{ARCHIVE}'
# Published in https://download.blender.org/release/Blender5.1/blender-5.1.2.sha256
SHA256 = 'aaccb355f50183979b698bcce7467103a76261b5fa59f4972295842662a285fb'
DIRECTORY = f'blender-{VERSION}-linux-x64'
MIRRORS = (
    URL,
    f'https://mirrors.dotsrc.org/blender/release/Blender5.1/{ARCHIVE}',
    f'https://ftp.halifax.rwth-aachen.de/blender/release/Blender5.1/{ARCHIVE}',
)
USER_AGENT = 'ComponentStudio/0.1 (pinned Blender runtime installer)'


def validate_download_url(url):
    parsed = urlsplit(url)
    if url not in MIRRORS or parsed.scheme != 'https' or parsed.username or parsed.password:
        raise ValueError('Blender download must use an approved release URL')


def download_archive(destination):
    """Try each official mirror once; checksum failures always stop installation."""
    errors = []
    for url in MIRRORS:
        validate_download_url(url)
        request = urllib.request.Request(url, headers={'User-Agent': USER_AGENT})
        try:
            with urllib.request.urlopen(request, timeout=120) as response, Path(destination).open('wb') as stream:
                validate_download_url(response.geturl())
                shutil.copyfileobj(response, stream, length=1024 * 1024)
        except (urllib.error.URLError, TimeoutError, ConnectionError, http.client.HTTPException) as error:
            Path(destination).unlink(missing_ok=True)
            errors.append(f'{urlsplit(url).hostname}: {type(error).__name__}')
            continue
        with Path(destination).open('rb') as stream:
            if hashlib.file_digest(stream, 'sha256').hexdigest() != SHA256:
                Path(destination).unlink(missing_ok=True)
                raise RuntimeError('Official Blender archive checksum mismatch')
        return url
    raise RuntimeError('All official Blender downloads failed: ' + '; '.join(errors))



def verified_extract(archive, directory, expected_sha256=SHA256):
    with Path(archive).open('rb') as stream:
        actual = hashlib.file_digest(stream, 'sha256').hexdigest()
    if actual != expected_sha256:
        raise RuntimeError('Official Blender archive checksum mismatch')
    with tarfile.open(archive, 'r:xz') as bundle:
        bundle.extractall(directory, filter='data')


def probe_blender(binary):
    """Render with the actual OIDN CPU kernel, rather than trusting build metadata."""
    expression = (
        'import bpy, _cycles, json; '
        f'assert bpy.app.version_string == "{VERSION}", bpy.app.version_string; '
        'assert getattr(_cycles, "with_openimagedenoise", False), "OpenImageDenoise unavailable"; '
        's=bpy.context.scene; s.render.engine="CYCLES"; s.cycles.device="CPU"; '
        's.cycles.samples=2; s.cycles.use_denoising=True; s.cycles.denoiser="OPENIMAGEDENOISE"; '
        's.view_settings.view_transform="AgX"; '
        's.render.resolution_x=32; s.render.resolution_y=32; s.render.resolution_percentage=100; '
        'bpy.ops.render.render(); '
        'print("STUDIO_BLENDER_PROBE="+json.dumps({"version":bpy.app.version_string, '
        '"openimagedenoise":True,"rendered":True,"view_transform":s.view_settings.view_transform}))'
    )
    result = subprocess.run([str(binary), '--background', '--factory-startup', '--python-exit-code', '1',
                             '--python-expr', expression], capture_output=True, text=True, timeout=180)
    if result.returncode:
        raise RuntimeError('Pinned Blender CPU denoising render probe failed: ' + result.stdout[-2000:]
                           + result.stderr[-1000:])
    markers = [line.split('=', 1)[1] for line in result.stdout.splitlines()
               if line.startswith('STUDIO_BLENDER_PROBE=')]
    if len(markers) != 1:
        raise RuntimeError('Pinned Blender did not return its render probe receipt')
    return json.loads(markers[0])


def prepare_blender(runtime):
    """Return a verified binary path; caller sets BLENDER_BIN in its runtime environment."""
    if platform.system() != 'Linux' or platform.machine() not in ('x86_64', 'AMD64'):
        raise RuntimeError('Pinned Blender installer requires Linux x86_64')
    runtime = Path(runtime).resolve()
    runtime.mkdir(parents=True, exist_ok=True)
    installed = runtime / DIRECTORY
    binary = installed / 'blender'
    receipt_path = installed / 'studio-install.json'
    if binary.is_file() and receipt_path.is_file():
        receipt = json.loads(receipt_path.read_text())
        if receipt.get('archive_sha256') != SHA256 or receipt.get('version') != VERSION:
            raise RuntimeError('Existing pinned Blender installation has an unexpected receipt')
        probe_blender(binary)
        return binary
    if installed.exists():
        raise RuntimeError('Incomplete Blender installation exists; remove it before reinstalling')
    with tempfile.TemporaryDirectory(prefix='blender-install-', dir=runtime) as temporary:
        staging = Path(temporary)
        archive = staging / ARCHIVE
        downloaded_from = download_archive(archive)
        verified_extract(archive, staging)
        candidate = staging / DIRECTORY
        probe = probe_blender(candidate / 'blender')
        receipt = {'version': VERSION, 'source': downloaded_from, 'archive_sha256': SHA256, 'probe': probe}
        (candidate / 'studio-install.json').write_text(json.dumps(receipt, indent=2))
        os.replace(candidate, installed)
    return binary


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--runtime', type=Path, default=Path('.runtime'))
    args = parser.parse_args()
    print(prepare_blender(args.runtime))


if __name__ == '__main__':
    main()