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: No API key set. Get a free key at https://tiingo.com or https://twelvedata.com\n");
            return;
        },
        else => {
            try cli.stderrPrint("Error fetching data.\n");
            return;
        },
    };
    defer allocator.free(result.candles);
    defer if (result.dividends) |d| zfin.Dividend.freeSlice(allocator, 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.fmtMoneyAbs(&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);
    }

    // -- Risk metrics --
    const tr = zfin.risk.trailingRisk(c);
    try printRiskTable(out, tr, 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) 1.0 else -1.0);
        } 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) 1.0 else -1.0);
                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", .{});
    }
}

pub fn printRiskTable(out: *std.Io.Writer, tr: zfin.risk.TrailingRisk, color: bool) !void {
    const risk_arr = [4]?zfin.risk.RiskMetrics{ tr.one_year, tr.three_year, tr.five_year, tr.ten_year };

    // Only show if at least one period has data
    var any = false;
    for (risk_arr) |r| {
        if (r != null) {
            any = true;
            break;
        }
    }
    if (!any) return;

    try cli.setBold(out, color);
    try out.print("\nRisk Metrics (monthly returns):\n", .{});
    try cli.reset(out, color);

    try cli.setFg(out, color, cli.CLR_MUTED);
    try out.print("{s:>22} {s:>14} {s:>14} {s:>14}\n", .{ "", "Volatility", "Sharpe", "Max DD" });
    try out.print("{s:->22} {s:->14} {s:->14} {s:->14}\n", .{ "", "", "", "" });
    try cli.reset(out, color);

    const labels = [4][]const u8{ "1-Year:", "3-Year:", "5-Year:", "10-Year:" };

    for (0..4) |i| {
        try out.print("  {s:<20}", .{labels[i]});
        if (risk_arr[i]) |rm| {
            try out.print(" {d:>12.1}%", .{rm.volatility * 100.0});
            try out.print(" {d:>13.2}", .{rm.sharpe});
            try cli.setFg(out, color, cli.CLR_NEGATIVE);
            try out.print(" {d:>12.1}%", .{rm.max_drawdown * 100.0});
            try cli.reset(out, color);
        } else {
            try cli.setFg(out, color, cli.CLR_MUTED);
            try out.print(" {s:>13} {s:>13} {s:>13}", .{ "—", "—", "—" });
            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);
}