"""画像専用の独立子プロセス。分離のGPU経路を利用しない。"""
import hashlib
import json
import os
from pathlib import Path
import sys
import warnings

os.environ.setdefault('OMP_NUM_THREADS', '1')
os.environ.setdefault('OPENBLAS_NUM_THREADS', '1')
os.environ.setdefault('CUDA_VISIBLE_DEVICES', '')

def write(folder, **values):
    p = folder/'progress.tmp'
    p.write_text(json.dumps(values), encoding='utf-8')
    p.replace(folder/'progress.json')

def model_file(name, expected):
    root = Path(os.environ.get('LUNARON_IMAGE_MODEL_DIR', Path(__file__).parent/'image_models'))
    p = root/name
    if hashlib.sha256(p.read_bytes()).hexdigest() != expected:
        raise ValueError('MODEL_INTEGRITY')
    return p

def open_image(p, mode):
    from PIL import Image, ImageOps
    Image.MAX_IMAGE_PIXELS = 16000000
    warnings.simplefilter('error', Image.DecompressionBombWarning)
    im = Image.open(p)
    if im.format not in {'JPEG', 'PNG', 'WEBP'} or getattr(im, 'n_frames', 1) != 1:
        raise ValueError('INVALID_IMAGE')
    if im.width*im.height > 16000000:
        raise ValueError('IMAGE_TOO_LARGE')
    im.load()
    im = ImageOps.exif_transpose(im).convert('RGBA')
    if mode == 'upscale' and (max(im.size) > 1024 or im.width*im.height > 1048576):
        raise ValueError('UPSCALE_LIMIT')
    return im

def background(folder, im):
    import numpy as np
    import onnxruntime as ort
    from PIL import Image, ImageChops
    path = model_file('birefnet-lite.onnx', '5600024376f572a557870a5eb0afb1e5961636bef4e1e22132025467d0f03333')
    options = ort.SessionOptions()
    options.intra_op_num_threads = 1
    options.inter_op_num_threads = 1
    session = ort.InferenceSession(str(path), sess_options=options, providers=['CPUExecutionProvider'])
    rgb = Image.new('RGB', im.size, 'white')
    rgb.paste(im, mask=im.getchannel('A'))
    arr = np.asarray(rgb.resize((1024, 1024), Image.Resampling.LANCZOS), dtype=np.float32)/255.0
    arr = ((arr-np.array([.485,.456,.406],dtype=np.float32))/np.array([.229,.224,.225],dtype=np.float32)).transpose(2,0,1)[None]
    write(folder, stage='processing')
    logits = session.run(None, {session.get_inputs()[0].name:arr})[0][0,0]
    mask = 1/(1+np.exp(-np.clip(logits,-30,30)))
    # 固定の確率を使い、背景がない画像をmin/max拡大で無理に切り抜かない。
    mask = Image.fromarray(np.uint8(np.clip(mask*255,0,255)), 'L')
    im.thumbnail((2048,2048), Image.Resampling.LANCZOS)
    mask = mask.resize(im.size, Image.Resampling.LANCZOS)
    im.putalpha(ImageChops.multiply(im.getchannel('A'), mask))
    write(folder, stage='saving')
    im.save(folder/'transparent.png', compress_level=4)
    white = Image.new('RGBA',im.size,'white')
    white.alpha_composite(im)
    white.convert('RGB').save(folder/'white.png', compress_level=4)
    return im

def upscale(folder, im):
    import numpy as np
    import torch
    from torch import nn
    from torch.nn import functional as F
    from PIL import Image
    torch.set_num_threads(1)
    torch.set_num_interop_threads(1)
    # Real-ESRGAN SRVGGNetCompact。著作権・BSD-3-Clauseはimage_modelsに同梱。
    class SRVGG(nn.Module):
        def __init__(self):
            super().__init__()
            self.body=nn.ModuleList([nn.Conv2d(3,64,3,1,1),nn.PReLU(64)])
            for _ in range(32):
                self.body.extend([nn.Conv2d(64,64,3,1,1),nn.PReLU(64)])
            self.body.append(nn.Conv2d(64,48,3,1,1))
            self.upsampler=nn.PixelShuffle(4)
        def forward(self,x):
            out=x
            for layer in self.body: out=layer(out)
            return self.upsampler(out)+F.interpolate(x,scale_factor=4,mode='nearest')
    p=model_file('realesr-general-x4v3.pth','8dc7edb9ac80ccdc30c3a5dca6616509367f05fbc184ad95b731f05bece96292')
    model=SRVGG().eval()
    checkpoint=torch.load(p,map_location='cpu',weights_only=True)
    model.load_state_dict(checkpoint['params_ema'] if 'params_ema' in checkpoint else checkpoint['params'],strict=True)
    rgb=Image.new('RGB',im.size,'white');rgb.paste(im,mask=im.getchannel('A'))
    arr=np.asarray(rgb,dtype=np.float32)/255
    w,h=im.size
    out=Image.new('RGB',(w*2,h*2))
    tiles=((w+127)//128)*((h+127)//128);done=0
    with torch.inference_mode():
        for y in range(0,h,128):
            for x in range(0,w,128):
                sx=max(0,x-20);sy=max(0,y-20);ex=min(w,x+148);ey=min(h,y+148)
                t=torch.from_numpy(arr[sy:ey,sx:ex].transpose(2,0,1).copy())[None]
                pred=model(t)[0].clamp(0,1).permute(1,2,0).numpy()
                crop=Image.fromarray(np.uint8(np.round(pred*255)))
                box=((x-sx)*4,(y-sy)*4,(min(w,x+128)-sx)*4,(min(h,y+128)-sy)*4)
                crop=crop.crop(box).resize((min(128,w-x)*2,min(128,h-y)*2),Image.Resampling.LANCZOS)
                out.paste(crop,(x*2,y*2));done+=1
                write(folder,stage='processing',done=done,total=tiles)
    # 透過の輪郭は元のアルファから補間し、黒背景として潰さない。
    if im.getchannel('A').getextrema() != (255,255):
        out=out.convert('RGBA');out.putalpha(im.getchannel('A').resize(out.size,Image.Resampling.LANCZOS))
    write(folder,stage='saving')
    out.save(folder/'enhanced.png',compress_level=4)
    return out

def run(folder):
    from PIL import Image
    if hasattr(os,'nice'): os.nice(10)
    config=json.loads((folder/'config.json').read_text(encoding='utf-8'))
    write(folder,stage='preparing')
    im=open_image(folder/'input',config['mode'])
    thumb=im.copy();thumb.thumbnail((480,480),Image.Resampling.LANCZOS)
    thumb.save(folder/'before.png')
    output=background(folder,im) if config['mode']=='background' else upscale(folder,im)
    preview=output.copy();preview.thumbnail((480,480),Image.Resampling.LANCZOS)
    preview.save(folder/'preview.png',compress_level=4)
    result={'stage':'done','width':output.width,'height':output.height}
    write(folder,**result)
    (folder/'input').unlink(missing_ok=True)
    return result

if __name__=='__main__':
    folder=Path(sys.argv[1]).resolve()
    try: run(folder)
    except Exception:
        write(folder,stage='error',error='IMAGE_PROCESSING_FAILED')
        (folder/'input').unlink(missing_ok=True)
        sys.exit(1)