"""Isolated A10 smoke test: prepare on CPU, then load and generate with Hub offline."""
import base64
import json
from pathlib import Path

import modal

from modal_app import ROOT, cache, image

validation_image = image.add_local_file(ROOT / 'verify_checkpoint.py', '/root/verify_checkpoint.py', copy=True)
for filename in ('modal_app.py', 'requirements.txt'):
    validation_image = validation_image.add_local_file(ROOT / filename, f'/root/{filename}', copy=True)
app = modal.App('socialtrait-flux3-checkpoint-validation')


@app.function(image=validation_image, cpu=2, memory=8192, timeout=1800,
              volumes={'/cache/huggingface': cache})
def prepare():
    from flux3_engine import Settings, download_weights
    from verify_checkpoint import verify_checkpoint
    settings = Settings.from_env()
    path = download_weights(settings)
    verification = verify_checkpoint(path, settings.quantization, check_hashes=True)
    cache.commit()
    return {'path': path, 'verification': verification}


@app.function(image=validation_image.env({'HF_HUB_OFFLINE': '1', 'TRANSFORMERS_OFFLINE': '1'}),
              gpu='A10', cpu=4, memory=32768, timeout=1200, max_containers=1,
              volumes={'/cache/huggingface': cache})
def generate_offline(path: str):
    from handler import EndpointHandler
    # Exercise the mounted-repository HF handler with network downloads disabled.
    handler = EndpointHandler(path)
    return handler({'inputs': 'A photographic portrait of an adult woman, natural face and eyes, bright background',
                    'width': 1024, 'height': 1024, 'seed': 42})


@app.local_entrypoint()
def main(output_dir: str = '/tmp/flux-selfcontained-validation'):
    prepared = prepare.remote()
    print('Verified selected Socialtrait weights:', prepared['verification'])
    result = generate_offline.remote(prepared['path'])
    output = Path(output_dir)
    output.mkdir(parents=True, exist_ok=True)
    (output / 'sample.jpg').write_bytes(base64.b64decode(result.pop('image')))
    result.update(status='success', hub_offline=True, verification=prepared['verification'])
    (output / 'results.json').write_text(json.dumps(result, indent=2) + '\n')
    print(json.dumps(result, indent=2))