from pathlib import Path

from mflux.callbacks.callback_manager import CallbackManager
from mflux.cli.defaults import defaults as ui_defaults
from mflux.cli.parser.parsers import CommandLineParser
from mflux.models.common.config import ModelConfig
from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator
from mflux.models.flux.variants.in_context.flux_in_context_fill import Flux1InContextFill
from mflux.models.flux.variants.in_context.utils.in_context_fill_util import FluxInContextFillUtil
from mflux.models.flux.variants.in_context.utils.in_context_loras import IC_EDIT_LORA_SCALE, get_ic_edit_lora_path
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException


def main():
    # 0. Parse command line arguments
    parser = CommandLineParser(description="Generate images using in-context editing.")
    parser.add_general_arguments()
    parser.add_model_arguments(require_model_arg=False)
    parser.add_lora_arguments()
    parser.add_image_generator_arguments(supports_metadata_config=False, require_prompt=False)
    parser.add_in_context_edit_arguments()
    parser.add_in_context_arguments()
    parser.add_output_arguments()
    args = parser.parse_args()

    # 0. Default to a higher guidance value for fill related tasks.
    if args.guidance is None:
        args.guidance = ui_defaults.DEFAULT_DEV_FILL_GUIDANCE

    # Auto-resize to optimal width for IC-Edit
    width, height = FluxInContextFillUtil.resize_for_ic_edit_optimal_width(args)

    # Build lora_paths: IC-Edit LoRA (required) + user-provided LoRAs
    lora_paths = [get_ic_edit_lora_path()]
    lora_scales = [IC_EDIT_LORA_SCALE]
    if args.lora_paths:
        lora_paths.extend(args.lora_paths)
        lora_scales.extend(args.lora_scales or [1.0] * len(args.lora_paths))

    # 1. Load the model with IC-Edit LoRA
    flux = Flux1InContextFill(
        model_config=ModelConfig.dev_fill(),
        quantize=args.quantize,
        model_path=args.model_path,
        lora_paths=lora_paths,
        lora_scales=lora_scales,
    )

    # 2. Register callbacks
    memory_saver = CallbackManager.register_callbacks(
        args=args,
        model=flux,
        latent_creator=FluxLatentCreator,
    )

    try:
        for seed in args.seed:
            # 3. Generate an image for each seed value
            image = flux.generate_image(
                seed=seed,
                prompt=FluxInContextFillUtil.get_effective_ic_edit_prompt(args),
                width=width,
                height=height,
                guidance=args.guidance,
                scheduler=args.scheduler,
                num_inference_steps=args.steps,
                left_image_path=args.reference_image,
                right_image_path=None,
            )

            # 4. Save the image(s)
            output_path = Path(args.output.format(seed=seed))
            image.get_right_half().save(path=output_path, export_json_metadata=args.metadata)
            if args.save_full_image:
                image.save(path=output_path.with_stem(output_path.stem + "_full"))

    except (StopImageGenerationException, PromptFileReadError) as exc:
        print(exc)
    finally:
        if memory_saver:
            print(memory_saver.memory_stats())


if __name__ == "__main__":
    main()
