"""mlx-music CLI."""

import argparse
import os
import sys
import time


def _print_status(msg, end="\n"):
    """Print a status message with flush for real-time output."""
    print(msg, end=end, flush=True)


def _format_time(seconds):
    """Format seconds to a display string."""
    return f"{seconds:.1f}s"


def cmd_generate(args):
    """Run the generate subcommand."""
    from mlx_music import __version__

    _print_status(f"mlx-music v{__version__} — ACE-Step 1.5")

    # --- Validate inputs ---
    if not args.caption or not args.caption.strip():
        print("Error: --caption is required and cannot be empty.", file=sys.stderr)
        print(
            '  Example: mlx-music generate --caption "A melodic piano ballad"',
            file=sys.stderr,
        )
        sys.exit(1)

    weights_dir = os.path.expanduser(args.weights_dir)
    if not os.path.isdir(weights_dir):
        print(f"Error: Weights directory not found: {weights_dir}", file=sys.stderr)
        print("", file=sys.stderr)
        print("  To get started, convert the ACE-Step weights first:", file=sys.stderr)
        print("    mlx-music convert --output ./weights", file=sys.stderr)
        print("", file=sys.stderr)
        print("  This downloads and converts the model to MLX format.", file=sys.stderr)
        sys.exit(1)

    # --- Load pipeline ---
    _print_status("Loading models...     ", end="")
    t0 = time.time()

    try:
        from mlx_music.models.acestep.pipeline import AceStepPipeline

        pipeline = AceStepPipeline(weights_dir=weights_dir)
        pipeline.load_models()
    except ImportError:
        print("FAILED", flush=True)
        print("", file=sys.stderr)
        print("Error: AceStepPipeline not available.", file=sys.stderr)
        print(
            "  The pipeline module may not be installed or built yet.", file=sys.stderr
        )
        print("  Try: mlx-music convert --output ./weights", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print("FAILED", flush=True)
        print("", file=sys.stderr)
        print(f"Error loading models: {e}", file=sys.stderr)
        print("", file=sys.stderr)
        print("  Make sure you have converted weights:", file=sys.stderr)
        print("    mlx-music convert --output ./weights", file=sys.stderr)
        sys.exit(1)

    load_time = time.time() - t0
    _print_status(f"[done, {_format_time(load_time)}]")

    # --- Generate ---
    _print_status(f"Generating ({args.steps} steps)... ", end="")
    t0 = time.time()

    def step_callback(step, total):
        _print_status(f"[step {step}/{total}] ", end="")

    try:
        output_path = pipeline.generate(
            caption=args.caption,
            lyrics=args.lyrics,
            duration=args.duration,
            steps=args.steps,
            cfg_scale=args.cfg_scale,
            shift=args.shift,
            seed=args.seed,
            output=args.output,
        )
    except Exception as e:
        print("FAILED", flush=True)
        print("", file=sys.stderr)
        print(f"Error during generation: {e}", file=sys.stderr)
        sys.exit(1)

    gen_time = time.time() - t0
    _print_status(f"[done, {_format_time(gen_time)}]")

    # --- Summary ---
    _print_status(f"Saved: {output_path} ({args.duration:.1f}s, 48kHz stereo)")


def cmd_repaint(args):
    """Run the repaint subcommand."""
    from mlx_music import __version__

    _print_status(f"mlx-music v{__version__} — ACE-Step 1.5 (repaint)")

    # --- Validate inputs ---
    if not args.caption or not args.caption.strip():
        print("Error: --caption is required and cannot be empty.", file=sys.stderr)
        sys.exit(1)

    if not os.path.isfile(args.src_audio):
        print(f"Error: Source audio not found: {args.src_audio}", file=sys.stderr)
        sys.exit(1)

    weights_dir = os.path.expanduser(args.weights_dir)
    if not os.path.isdir(weights_dir):
        print(f"Error: Weights directory not found: {weights_dir}", file=sys.stderr)
        sys.exit(1)

    # --- Load pipeline ---
    _print_status("Loading models...     ", end="")
    t0 = time.time()

    try:
        from mlx_music.models.acestep.pipeline import AceStepPipeline

        pipeline = AceStepPipeline(weights_dir=weights_dir)
        pipeline.load_models()
    except Exception as e:
        print("FAILED", flush=True)
        print(f"Error loading models: {e}", file=sys.stderr)
        sys.exit(1)

    load_time = time.time() - t0
    _print_status(f"[done, {_format_time(load_time)}]")

    # --- Repaint ---
    region = ""
    if args.repaint_start > 0 or args.repaint_end >= 0:
        end_str = f"{args.repaint_end:.1f}s" if args.repaint_end >= 0 else "end"
        region = f", region {args.repaint_start:.1f}s-{end_str}"
    _print_status(f"Repainting ({args.steps} steps{region})... ", end="")
    t0 = time.time()

    def step_callback(step, total):
        _print_status(f"[step {step}/{total}] ", end="")

    try:
        output_path = pipeline.generate_repaint(
            src_audio=args.src_audio,
            caption=args.caption,
            lyrics=args.lyrics,
            repaint_start=args.repaint_start,
            repaint_end=args.repaint_end,
            steps=args.steps,
            cfg_scale=args.cfg_scale,
            shift=args.shift,
            seed=args.seed,
            output=args.output,
        )
    except Exception as e:
        print("FAILED", flush=True)
        print(f"Error during repaint: {e}", file=sys.stderr)
        sys.exit(1)

    gen_time = time.time() - t0
    _print_status(f"[done, {_format_time(gen_time)}]")
    _print_status(f"Saved: {output_path} (48kHz stereo)")


def cmd_convert(args):
    """Run the convert subcommand."""
    from mlx_music.models.acestep.convert import convert_all

    convert_all(args.output)


def main():
    parser = argparse.ArgumentParser(
        description="mlx-music: Music generation on Apple Silicon with MLX"
    )
    parser.add_argument("--version", action="store_true", help="Show version")
    subparsers = parser.add_subparsers(dest="command")

    # --- convert ---
    convert_parser = subparsers.add_parser(
        "convert", help="Convert ACE-Step weights to MLX format"
    )
    convert_parser.add_argument(
        "--output",
        type=str,
        default="./weights",
        help="Output directory for converted weights (default: ./weights)",
    )

    # --- generate ---
    generate_parser = subparsers.add_parser(
        "generate", help="Generate music from a text description"
    )
    generate_parser.add_argument(
        "--model",
        type=str,
        default="acestep",
        choices=["acestep"],
        help="Model backend (default: acestep)",
    )
    generate_parser.add_argument(
        "--caption",
        type=str,
        required=True,
        help="Text description of the music to generate",
    )
    generate_parser.add_argument(
        "--lyrics",
        type=str,
        default="[Instrumental]",
        help="Song lyrics with section markers (default: [Instrumental])",
    )
    generate_parser.add_argument(
        "--duration",
        type=float,
        default=30.0,
        help="Duration in seconds (default: 30.0)",
    )
    generate_parser.add_argument(
        "--steps",
        type=int,
        default=8,
        help="Number of denoising steps (default: 8 for turbo)",
    )
    generate_parser.add_argument(
        "--cfg-scale",
        type=float,
        default=1.0,
        help="CFG guidance scale (default: 1.0, turbo has it baked in)",
    )
    generate_parser.add_argument(
        "--shift",
        type=float,
        default=3.0,
        help="Time schedule shift (default: 3.0 for turbo)",
    )
    generate_parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed for reproducibility (optional)",
    )
    generate_parser.add_argument(
        "--output",
        type=str,
        default="output.wav",
        help="Output WAV file path (default: output.wav)",
    )
    generate_parser.add_argument(
        "--weights-dir",
        type=str,
        default="./weights",
        help="Directory containing converted weights (default: ./weights)",
    )

    # --- repaint ---
    repaint_parser = subparsers.add_parser(
        "repaint", help="Regenerate a region of existing audio guided by a text prompt"
    )
    repaint_parser.add_argument(
        "--src-audio",
        type=str,
        required=True,
        help="Path to source audio file (WAV, FLAC, OGG, etc.)",
    )
    repaint_parser.add_argument(
        "--caption",
        type=str,
        required=True,
        help="Text description of the music to generate in the repainted region",
    )
    repaint_parser.add_argument(
        "--lyrics",
        type=str,
        default="[Instrumental]",
        help="Song lyrics with section markers (default: [Instrumental])",
    )
    repaint_parser.add_argument(
        "--repaint-start",
        type=float,
        default=0.0,
        help="Start of region to regenerate in seconds (default: 0.0)",
    )
    repaint_parser.add_argument(
        "--repaint-end",
        type=float,
        default=-1.0,
        help="End of region to regenerate in seconds (-1 = full duration, default: -1)",
    )
    repaint_parser.add_argument(
        "--steps",
        type=int,
        default=8,
        help="Number of denoising steps (default: 8 for turbo)",
    )
    repaint_parser.add_argument(
        "--cfg-scale",
        type=float,
        default=1.0,
        help="CFG guidance scale (default: 1.0)",
    )
    repaint_parser.add_argument(
        "--shift",
        type=float,
        default=3.0,
        help="Time schedule shift (default: 3.0 for turbo)",
    )
    repaint_parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed for reproducibility (optional)",
    )
    repaint_parser.add_argument(
        "--output",
        type=str,
        default="output.wav",
        help="Output WAV file path (default: output.wav)",
    )
    repaint_parser.add_argument(
        "--weights-dir",
        type=str,
        default="./weights",
        help="Directory containing converted weights (default: ./weights)",
    )

    args = parser.parse_args()

    if args.version:
        from mlx_music import __version__

        print(f"mlx-music {__version__}")
    elif args.command == "convert":
        cmd_convert(args)
    elif args.command == "generate":
        cmd_generate(args)
    elif args.command == "repaint":
        cmd_repaint(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()
