"""Isolated Qwen Image 2.1 research worker; the 2511 handler stays unchanged."""
import io
import json
import os
import time
import traceback
import uuid
from pathlib import Path

os.environ['HF_HUB_CACHE'] = '/runpod-volume/huggingface-cache/hub'
os.environ['HF_HUB_OFFLINE'] = '1'
os.environ['TRANSFORMERS_OFFLINE'] = '1'

import requests
from PIL import Image

_pipe = None
_bucket = None


def sampling(job):
    allowed = {'job_id', 'image_url', 'image_urls', 'prompt', 'negative_prompt',
               'seed', 'size', 'num_inference_steps', 'true_cfg_scale'}
    if set(job) - allowed:
        raise ValueError('Unsupported job fields: ' + ', '.join(sorted(set(job) - allowed)))
    width, height = map(int, job.get('size', '1024*1536').split('*'))
    steps = int(job.get('num_inference_steps', 40))
    cfg = float(job.get('true_cfg_scale', 1.0))
    if not (64 <= width <= 4096 and 64 <= height <= 4096 and
            width % 16 == height % 16 == 0 and 1 <= steps <= 60 and 1 <= cfg <= 5):
        raise ValueError('Invalid resolution or sampling parameters')
    if not job.get('prompt', '').strip():
        raise ValueError('An explicit editing prompt is required')
    return {'prompt': job['prompt'], 'negative_prompt': job.get('negative_prompt'),
            'width': width, 'height': height, 'num_inference_steps': steps,
            'true_cfg_scale': cfg, 'use_kv_cache': True,
            'output_resolution': 1024}


def cached_model_path():
    folder = Path(os.environ['HF_HUB_CACHE']) / (
        'models--' + os.environ['QWEN21_MODEL_REPO'].replace('/', '--')) / (
        'snapshots/' + os.environ['QWEN21_MODEL_REVISION'])
    if not (folder / 'model_index.json').is_file():
        raise FileNotFoundError('Pinned model is missing from RunPod model cache: ' + str(folder))
    return str(folder)


def pipeline():
    global _pipe
    if _pipe is None:
        import torch
        from diffusers import QwenImage21Pipeline
        _pipe = QwenImage21Pipeline.from_pretrained(
            cached_model_path(), local_files_only=True,
            torch_dtype=torch.bfloat16)
        _pipe.enable_model_cpu_offload()
    return _pipe


def upload(photo):
    global _bucket
    if _bucket is None:
        from google.cloud import storage
        from google.oauth2 import service_account
        raw = os.environ['JSON_PATH_KEY']
        try:
            credentials = service_account.Credentials.from_service_account_info(json.loads(raw))
        except json.JSONDecodeError:
            credentials = service_account.Credentials.from_service_account_file(raw)
        _bucket = storage.Client(credentials=credentials).bucket(os.environ.get('GCS_BUCKET_NAME', 'shoupapp'))
    buf = io.BytesIO()
    photo.save(buf, 'PNG')
    blob = _bucket.blob('research/qwen21/' + uuid.uuid4().hex + '.png')
    blob.upload_from_string(buf.getvalue(), content_type='image/png')
    return blob.public_url


def handler(event):
    data = event['input']
    if data.get('task') != 'edit' or not 1 <= len(data.get('jobs', [])) <= 4:
        raise ValueError('Research worker accepts one to four explicit edit jobs')
    results = []
    for job in data['jobs']:
        try:
            args = sampling(job)
            urls = job.get('image_urls') or [job.get('image_url')]
            if not 1 <= len(urls) <= 10 or not all(urls):
                raise ValueError('One to ten source image URLs are required')
            images = []
            for url in urls:
                response = requests.get(url, timeout=60)
                response.raise_for_status()
                images.append(Image.open(io.BytesIO(response.content)).convert('RGBA'))
            import torch
            pipe = pipeline()
            torch.cuda.synchronize()
            torch.cuda.reset_peak_memory_stats()
            started = time.monotonic()
            photo = pipe(image=images, generator=torch.Generator('cuda').manual_seed(
                int(job.get('seed', 20261040))), **args).images[0]
            torch.cuda.synchronize()
            inference_ms = round((time.monotonic() - started) * 1000)
            results.append({'job_id': job.get('job_id'), 'ok': True,
                            'output_url': upload(photo.convert('RGB')),
                            'native_alpha_range': photo.getchannel('A').getextrema() if photo.mode == 'RGBA' else None,
                            'inference_ms': inference_ms,
                            'denoising_steps': len(pipe.scheduler.timesteps),
                            'gpu': torch.cuda.get_device_name(),
                            'peak_reserved_gib': round(torch.cuda.max_memory_reserved() / 1024**3, 2),
                            'model_revision': os.environ['QWEN21_MODEL_REVISION']})
        except Exception as exc:
            detail = traceback.format_exc()
            for key in ('HF_TOKEN', 'JSON_PATH_KEY'):
                if os.environ.get(key):
                    detail = detail.replace(os.environ[key], '[REDACTED]')
            print('Research job failed:', type(exc).__name__, detail, flush=True)
            results.append({'job_id': job.get('job_id'), 'ok': False,
                            'error': type(exc).__name__ + ': edit failed; inspect private worker logs'})
    return {'results': results}


if __name__ == '__main__':
    import runpod
    runpod.serverless.start({'handler': handler})