"""Explicit local Quasar inference. No model/codec download or tool execution."""
from pathlib import Path
import argparse
import os
import sys

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--checkpoint', type=Path, default=ROOT)
    parser.add_argument('--task', choices=('text', 'image', 'video'), required=True)
    parser.add_argument('--prompt', required=True)
    parser.add_argument('--device', choices=('cpu', 'cuda'), default='cuda')
    parser.add_argument('--output', type=Path)
    parser.add_argument('--resolution', type=int, choices=(64, 128), default=128)
    parser.add_argument('--frames', type=int, default=8)
    parser.add_argument('--seed', type=int, default=20261021)
    parser.add_argument('--max-new-tokens', type=int, default=64)
    args = parser.parse_args()
    if args.device == 'cpu':
        os.environ['CUDA_VISIBLE_DEVICES'] = ''
        os.environ.setdefault('TORCHINDUCTOR_TRITON_DISABLE_DEVICE_DETECTION', '1')
        # These FP32/BF16 weights do not use the optional GPU quantizer.
        sys.modules['torchao'] = None
    import torch
    torch.set_num_threads(2)
    if args.device == 'cuda':
        from quasar.training.hardware import check_gpu
        check_gpu({'require_cuda': True, 'exclusive_gpu': True, 'min_free_vram_gb': 4})
    from quasar.inference.extended import ExtendedGenerator
    from quasar.inference.decoding import Sampling
    generator = ExtendedGenerator(args.checkpoint.resolve(), args.device)
    sampling = Sampling(temperature=.5, top_k=64, top_p=.9, seed=args.seed, max_new_tokens=args.max_new_tokens)
    if args.task == 'text':
        result = generator.generate_text(args.prompt, sampling)
        if args.output:
            args.output.parent.mkdir(parents=True, exist_ok=True)
            args.output.write_text(result, encoding='utf-8')
        print(result)
        return
    from PIL import Image
    output = args.output or Path('generated.png' if args.task == 'image' else 'generated.gif')
    output.parent.mkdir(parents=True, exist_ok=True)
    if args.task == 'image':
        pixels = generator.generate_image(args.prompt, sampling, resolution=args.resolution)
        Image.fromarray(pixels).save(output)
    else:
        if output.suffix.lower() != '.gif':
            parser.error('Video preview output must use .gif')
        frames = generator.generate_video(args.prompt, args.frames, sampling)
        images = [Image.fromarray(frame) for frame in frames]
        images[0].save(output, save_all=True, append_images=images[1:], duration=250, loop=0, disposal=2)
    print(output.resolve())


if __name__ == '__main__':
    main()