const std = @import("std");
const zfin = @import("../root.zig");
const cli = @import("common.zig");
const fmt = cli.fmt;

pub fn run(allocator: std.mem.Allocator, svc: *zfin.DataService, symbol: []const u8, color: bool, out: *std.Io.Writer) !void {
    const result = svc.getTrailingReturns(symbol) catch |err| switch (err) {
        zfin.DataError.NoApiKey => {
            try cli.stderrPrint("Error: TWELVEDATA_API_KEY not set. Get a free key at https://twelvedata.com\n");
            return;
        },
        else => {
            try cli.stderrPrint("Error fetching data.\n");
            return;
        },
    };
    defer allocator.free(result.candles);
    defer if (result.dividends) |d| allocator.free(d);

    if (result.source == .cached) try cli.stderrPrint("(using cached data)\n");

    const c = result.candles;
    const end_date = c[c.len - 1].date;
    const today = fmt.todayDate();
    const month_end = today.lastDayOfPriorMonth();

    try cli.setBold(out, color);
    try out.print("\nTrailing Returns for {s}\n", .{symbol});
    try cli.reset(out, color);
    try out.print("========================================\n", .{});
    try cli.setFg(out, color, cli.CLR_MUTED);
    try out.print("Data points: {d} (", .{c.len});
    {
        var db: [10]u8 = undefined;
        try out.print("{s}", .{c[0].date.format(&db)});
    }
    try out.print(" to ", .{});
    {
        var db: [10]u8 = undefined;
        try out.print("{s}", .{end_date.format(&db)});
    }
    try cli.reset(out, color);
    var close_buf: [24]u8 = undefined;
    try out.print(")\nLatest close: {s}\n", .{fmt.fmtMoney(&close_buf, c[c.len - 1].close)});

    const has_divs = result.asof_total != null;

    // -- As-of-date returns --
    {
        var db: [10]u8 = undefined;
        try cli.setBold(out, color);
        try out.print("\nAs-of {s}:\n", .{end_date.format(&db)});
        try cli.reset(out, color);
    }
    try printReturnsTable(out, result.asof_price, if (has_divs) result.asof_total else null, color);

    // -- Month-end returns --
    {
        var db: [10]u8 = undefined;
        try cli.setBold(out, color);
        try out.print("\nMonth-end ({s}):\n", .{month_end.format(&db)});
        try cli.reset(out, color);
    }
    try printReturnsTable(out, result.me_price, if (has_divs) result.me_total else null, color);

    if (!has_divs) {
        try cli.setFg(out, color, cli.CLR_MUTED);
        try out.print("\nSet POLYGON_API_KEY for total returns with dividend reinvestment.\n", .{});
        try cli.reset(out, color);
    }
    try out.print("\n", .{});
}

pub fn printReturnsTable(
    out: *std.Io.Writer,
    price: zfin.performance.TrailingReturns,
    total: ?zfin.performance.TrailingReturns,
    color: bool,
) !void {
    const has_total = total != null;

    try cli.setFg(out, color, cli.CLR_MUTED);
    if (has_total) {
        try out.print("{s:>22} {s:>14} {s:>14}\n", .{ "", "Price Only", "Total Return" });
        try out.print("{s:->22} {s:->14} {s:->14}\n", .{ "", "", "" });
    } else {
        try out.print("{s:>22} {s:>14}\n", .{ "", "Price Only" });
        try out.print("{s:->22} {s:->14}\n", .{ "", "" });
    }
    try cli.reset(out, color);

    const periods = [_]struct { label: []const u8, years: u16 }{
        .{ .label = "1-Year Return:", .years = 1 },
        .{ .label = "3-Year Return:", .years = 3 },
        .{ .label = "5-Year Return:", .years = 5 },
        .{ .label = "10-Year Return:", .years = 10 },
    };

    const price_arr = [_]?zfin.performance.PerformanceResult{
        price.one_year, price.three_year, price.five_year, price.ten_year,
    };

    const total_arr: [4]?zfin.performance.PerformanceResult = if (total) |t|
        .{ t.one_year, t.three_year, t.five_year, t.ten_year }
    else
        .{ null, null, null, null };

    for (periods, 0..) |period, i| {
        try out.print("  {s:<20}", .{period.label});

        var price_buf: [32]u8 = undefined;
        var total_buf: [32]u8 = undefined;
        const row = fmt.fmtReturnsRow(
            &price_buf,
            &total_buf,
            price_arr[i],
            if (has_total) total_arr[i] else null,
            period.years > 1,
        );

        if (price_arr[i] != null) {
            try cli.setGainLoss(out, color, if (row.price_positive) @as(f64, 1) else @as(f64, -1));
        } else {
            try cli.setFg(out, color, cli.CLR_MUTED);
        }
        try out.print(" {s:>13}", .{row.price_str});
        try cli.reset(out, color);

        if (has_total) {
            if (row.total_str) |ts| {
                try cli.setGainLoss(out, color, if (row.price_positive) @as(f64, 1) else @as(f64, -1));
                try out.print(" {s:>13}", .{ts});
                try cli.reset(out, color);
            } else {
                try cli.setFg(out, color, cli.CLR_MUTED);
                try out.print(" {s:>13}", .{"N/A"});
                try cli.reset(out, color);
            }
        }

        if (row.suffix.len > 0) {
            try cli.setFg(out, color, cli.CLR_MUTED);
            try out.print("{s}", .{row.suffix});
            try cli.reset(out, color);
        }
        try out.print("\n", .{});
    }
}

// ── Tests ────────────────────────────────────────────────────

test "printReturnsTable price-only with no data" {
    var buf: [4096]u8 = undefined;
    var w: std.Io.Writer = .fixed(&buf);
    const empty: zfin.performance.TrailingReturns = .{
        .one_year = null,
        .three_year = null,
        .five_year = null,
        .ten_year = null,
    };
    try printReturnsTable(&w, empty, null, false);
    const out = w.buffered();
    // Should contain header and N/A for all periods
    try std.testing.expect(std.mem.indexOf(u8, out, "Price Only") != null);
    try std.testing.expect(std.mem.indexOf(u8, out, "N/A") != null);
    try std.testing.expect(std.mem.indexOf(u8, out, "1-Year") != null);
    try std.testing.expect(std.mem.indexOf(u8, out, "10-Year") != null);
}

test "printReturnsTable price-only no ANSI without color" {
    var buf: [4096]u8 = undefined;
    var w: std.Io.Writer = .fixed(&buf);
    const empty: zfin.performance.TrailingReturns = .{
        .one_year = null,
        .three_year = null,
        .five_year = null,
        .ten_year = null,
    };
    try printReturnsTable(&w, empty, null, false);
    const out = w.buffered();
    // No ANSI escape sequences when color=false
    try std.testing.expect(std.mem.indexOf(u8, out, "\x1b[") == null);
}

test "printReturnsTable with total return columns" {
    var buf: [4096]u8 = undefined;
    var w: std.Io.Writer = .fixed(&buf);
    const empty: zfin.performance.TrailingReturns = .{
        .one_year = null,
        .three_year = null,
        .five_year = null,
        .ten_year = null,
    };
    try printReturnsTable(&w, empty, empty, false);
    const out = w.buffered();
    // Should contain both column headers
    try std.testing.expect(std.mem.indexOf(u8, out, "Price Only") != null);
    try std.testing.expect(std.mem.indexOf(u8, out, "Total Return") != null);
}

test "printReturnsTable with actual returns" {
    var buf: [4096]u8 = undefined;
    var w: std.Io.Writer = .fixed(&buf);
    const returns: zfin.performance.TrailingReturns = .{
        .one_year = .{ .total_return = 0.15, .annualized_return = null, .from = .{ .days = 0 }, .to = .{ .days = 365 } },
        .three_year = null,
        .five_year = null,
        .ten_year = null,
    };
    try printReturnsTable(&w, returns, null, false);
    const out = w.buffered();
    // 1-Year should show a value, not N/A
    // Check that the line with "1-Year" does NOT have N/A right after it
    // (crude check: the output should have fewer N/A occurrences than with all nulls)
    try std.testing.expect(std.mem.indexOf(u8, out, "1-Year") != null);
    // 3-year should still show N/A
    try std.testing.expect(std.mem.indexOf(u8, out, "ann.") != null);
}