const std = @import("std");
const Allocator = std.mem.Allocator;
const Io = std.Io;

const zignal = @import("zignal");

const args = @import("args.zig");
const common = @import("common.zig");
const display = @import("display.zig");

const Args = struct {
    display: bool = false,
    width: ?u32 = null,
    height: ?u32 = null,
    protocol: ?display.ProtocolTag = null,

    pub const meta = .{
        .display = .{ .help = "Display the result in the terminal", .short = 'd' },
        .width = .{ .help = "Width of each sub-image", .metavar = "N" },
        .height = .{ .help = "Height of each sub-image", .metavar = "N" },
        .protocol = .{ .help = display.protocol_help, .metavar = "p" },
    };
};

pub const description = "Apply Feature Distribution Matching (style transfer) from target to source image.\nIf output is omitted, the result is displayed in the terminal.";

pub const help = args.generateHelp(
    Args,
    "zignal fdm <source> <target> [output] [options]",
    description,
);

pub fn run(io: Io, writer: *Io.Writer, gpa: Allocator, iterator: *std.process.Args.Iterator) !void {
    const parsed = try args.parse(Args, gpa, iterator);
    defer parsed.deinit(gpa);

    if (parsed.help or parsed.positionals.len < 2 or parsed.positionals.len > 3) {
        try args.printHelp(writer, help);
        return;
    }

    const source_path = parsed.positionals[0];
    const target_path = parsed.positionals[1];
    const output_path = if (parsed.positionals.len == 3) parsed.positionals[2] else null;

    const should_display = parsed.options.display or output_path == null;

    const Pixel = zignal.Rgb(u8);

    std.log.debug("loading source image: {s}", .{source_path});
    var source_img: zignal.Image(Pixel) = try .load(io, gpa, source_path);
    defer source_img.deinit(gpa);

    // FDM mutates `source_img` in place; keep an original copy for the side-by-side preview.
    var original_source: ?zignal.Image(Pixel) = null;
    if (should_display) {
        original_source = try source_img.dupe(gpa);
    }
    defer if (original_source) |*img| img.deinit(gpa);

    std.log.debug("loading target image: {s}", .{target_path});
    var target_img: zignal.Image(Pixel) = try .load(io, gpa, target_path);
    defer target_img.deinit(gpa);

    var fdm: zignal.FeatureDistributionMatching(Pixel) = .init(gpa);
    defer fdm.deinit();

    std.log.debug("applying fdm style transfer...", .{});

    const timer = common.Timer.begin(io);
    try fdm.match(source_img, target_img);
    timer.logElapsed("fdm");

    if (output_path) |out_path| {
        std.log.info("saving result to {s}...", .{out_path});
        try source_img.save(io, gpa, out_path);
    }

    if (should_display) {
        const images = [_]zignal.Image(Pixel){ original_source.?, target_img, source_img };

        var canvas = try display.createHorizontalComposite(
            Pixel,
            gpa,
            &images,
            parsed.options.width,
            parsed.options.height,
        );
        defer canvas.deinit(gpa);

        const format = display.resolveDisplayFormat(parsed.options.protocol, null, null);
        try display.displayCanvas(io, writer, &canvas, format);
    }
}
