import argparse
import shutil
from pathlib import Path

from editing_DiTs.image_generators.flux_image_generator import (
    FluxImageGenerator,
)

DEFAULT_EMBEDDING_NAME = "embedding.pt"
DEFAULT_IMAGE_NAME = "output.png"
DEFAULT_OUTPUT_DIR = Path("results/images/")


def main() -> None:
    parser = argparse.ArgumentParser(description="Flux Inference Script")
    parser.add_argument(
        "-p",
        "--prompt",
        type=str,
        help="The prompt to generate the image either a string or a path to a prompt embedding.",
    )
    parser.add_argument(
        "-o",
        "--output-dir",
        type=Path,
        default=DEFAULT_OUTPUT_DIR,
        help=f"Directory where the output subdirectory embeddings and images will be created and outputs saved to. Defaults to {DEFAULT_OUTPUT_DIR}.",
    )
    parser.add_argument(
        "-e",
        "--embedding-name",
        type=str,
        default=DEFAULT_EMBEDDING_NAME,
        help=f"Name of the output embedding file; should end in .pt. Defaults to {DEFAULT_IMAGE_NAME}.",
    )
    parser.add_argument(
        "-s",
        "--schedule",
        type=Path,
        help="Path to the schedule JSON file (containing DiT and optionally caching schedule).",
        default=None,
    )
    parser.add_argument(
        "-n",
        "--num-images",
        type=int,
        default=10,
        help="Number of images to generate for each prompt.",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=256,
        help="Height of the generated image.",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=256,
        help="Width of the generated image.",
    )
    parser.add_argument(
        "--guidance-scale",
        type=float,
        default=3.5,
        help="Guidance scale for the image generation.",
    )
    args = parser.parse_args()

    if not args.output_dir.exists():
        print(f"Output directory {args.output_dir} not found. Creating.")
        args.output_dir.mkdir(parents=True)

    if args.schedule is not None and not args.schedule.exists():
        print(f"Schedule file {args.schedule} not found. Exiting.")
        return

    prompt_dir = args.output_dir / "embeddings"
    images_dir = args.output_dir / "images"
    prompt_dir.mkdir(parents=True, exist_ok=True)
    images_dir.mkdir(parents=True, exist_ok=True)

    image_generator = FluxImageGenerator(schedule_path=args.schedule)

    if args.prompt is None or not Path(args.prompt).exists():
        prompt_file = Path("embedding.pt")
        print(
            f"Prompt file not found. Generating to {prompt_dir / prompt_file}"
        )

        image_generator.encode_and_save_prompts(
            {prompt_file.name: args.prompt}, prompt_dir, free_after=True
        )
        print("Prompt file saved.")
    else:
        print(f"Copying prompt file {args.prompt} to {prompt_dir}")
        try:
            shutil.copy(args.prompt, prompt_dir)
        except shutil.SameFileError:
            print("Prompt file already in embeddings directory. Continuing.")

    image_generator.generate_from_saved_prompts(
        prompt_dir,
        images_dir,
        batch_size=1,
        images_per_prompt=args.num_images,
        free_after=True,
        height=args.height,
        width=args.width,
        guidance_scale=args.guidance_scale,
    )


if __name__ == "__main__":
    print("Running pixart_inference.py")
    main()
