const std = @import("std");
const Date = @import("../models/date.zig").Date;
const Candle = @import("../models/candle.zig").Candle;
const Dividend = @import("../models/dividend.zig").Dividend;

/// Minimum holding period (in years) before annualizing returns.
/// Set below 1.0 to handle trading-day snap (e.g. a "1-year" lookback
/// that lands on 362 days due to weekends).
const min_annualize_years = 0.95;

/// Compute CAGR from a total return over a given number of years.
/// Returns null for periods shorter than `min_annualize_years` where
/// extrapolating to a full year would be misleading.
inline fn annualizedReturn(total: f64, years: f64) ?f64 {
    if (years < min_annualize_years) return null;
    return std.math.pow(f64, 1.0 + total, 1.0 / years) - 1.0;
}

/// Performance calculation results, Morningstar-style.
pub const PerformanceResult = struct {
    /// Total return over the period (e.g., 0.25 = 25%)
    total_return: f64,
    /// Annualized return (for periods > 1 year)
    annualized_return: ?f64,
    /// Start date used
    from: Date,
    /// End date used
    to: Date,
};

/// Compute total return from adjusted close prices.
/// Candles must be sorted by date ascending.
/// `from` snaps forward (first trading day on/after), `to` snaps backward.
pub fn totalReturnFromAdjClose(candles: []const Candle, from: Date, to: Date) ?PerformanceResult {
    return totalReturnFromAdjCloseSnap(candles, from, to, .forward);
}

/// Same as totalReturnFromAdjClose but both dates snap backward
/// (last trading day on or before). Used for month-end methodology where
/// both from and to represent month-end reference dates.
fn totalReturnFromAdjCloseBackward(candles: []const Candle, from: Date, to: Date) ?PerformanceResult {
    return totalReturnFromAdjCloseSnap(candles, from, to, .backward);
}

fn totalReturnFromAdjCloseSnap(candles: []const Candle, from: Date, to: Date, start_dir: SearchDirection) ?PerformanceResult {
    const start = findNearestCandle(candles, from, start_dir) orelse return null;
    const end = findNearestCandle(candles, to, .backward) orelse return null;

    if (start.adj_close == 0) return null;

    const total = (end.adj_close / start.adj_close) - 1.0;
    const years = Date.yearsBetween(start.date, end.date);

    return .{
        .total_return = total,
        .annualized_return = annualizedReturn(total, years),
        .from = start.date,
        .to = end.date,
    };
}

/// Compute total return with manual dividend reinvestment.
/// Uses raw close prices and dividend records independently.
/// Candles and dividends must be sorted by date ascending.
/// `from` snaps forward, `to` snaps backward.
pub fn totalReturnWithDividends(
    candles: []const Candle,
    dividends: []const Dividend,
    from: Date,
    to: Date,
) ?PerformanceResult {
    return totalReturnWithDividendsSnap(candles, dividends, from, to, .forward);
}

/// Same as totalReturnWithDividends but both dates snap backward.
fn totalReturnWithDividendsBackward(
    candles: []const Candle,
    dividends: []const Dividend,
    from: Date,
    to: Date,
) ?PerformanceResult {
    return totalReturnWithDividendsSnap(candles, dividends, from, to, .backward);
}

fn totalReturnWithDividendsSnap(
    candles: []const Candle,
    dividends: []const Dividend,
    from: Date,
    to: Date,
    start_dir: SearchDirection,
) ?PerformanceResult {
    const start = findNearestCandle(candles, from, start_dir) orelse return null;
    const end = findNearestCandle(candles, to, .backward) orelse return null;

    if (start.close == 0) return null;

    // Simulate: start with 1 share, reinvest dividends at ex-date close
    var shares: f64 = 1.0;

    for (dividends) |div| {
        if (div.ex_date.lessThan(start.date)) continue;
        if (end.date.lessThan(div.ex_date)) break;

        // Find close price on or near the ex-date
        const price_candle = findNearestCandle(candles, div.ex_date, .backward) orelse continue;
        if (price_candle.close > 0) {
            shares += (div.amount * shares) / price_candle.close;
        }
    }

    const final_value = shares * end.close;
    const total = (final_value / start.close) - 1.0;
    const years = Date.yearsBetween(start.date, end.date);

    return .{
        .total_return = total,
        .annualized_return = annualizedReturn(total, years),
        .from = start.date,
        .to = end.date,
    };
}

/// Convenience: compute 1yr, 3yr, 5yr, 10yr trailing returns from adjusted close.
/// Uses the last available date as the endpoint.
pub const TrailingReturns = struct {
    one_year: ?PerformanceResult = null,
    three_year: ?PerformanceResult = null,
    five_year: ?PerformanceResult = null,
    ten_year: ?PerformanceResult = null,
};

/// Trailing returns from exact calendar date N years ago to latest candle date.
/// Start dates snap forward to the next trading day (e.g., weekend → Monday).
pub fn trailingReturns(candles: []const Candle) TrailingReturns {
    if (candles.len == 0) return .{};

    const end_date = candles[candles.len - 1].date;

    return .{
        .one_year = totalReturnFromAdjClose(candles, end_date.subtractYears(1), end_date),
        .three_year = totalReturnFromAdjClose(candles, end_date.subtractYears(3), end_date),
        .five_year = totalReturnFromAdjClose(candles, end_date.subtractYears(5), end_date),
        .ten_year = totalReturnFromAdjClose(candles, end_date.subtractYears(10), end_date),
    };
}

/// Same as trailingReturns but with dividend reinvestment.
pub fn trailingReturnsWithDividends(
    candles: []const Candle,
    dividends: []const Dividend,
) TrailingReturns {
    if (candles.len == 0) return .{};

    const end_date = candles[candles.len - 1].date;

    return .{
        .one_year = totalReturnWithDividends(candles, dividends, end_date.subtractYears(1), end_date),
        .three_year = totalReturnWithDividends(candles, dividends, end_date.subtractYears(3), end_date),
        .five_year = totalReturnWithDividends(candles, dividends, end_date.subtractYears(5), end_date),
        .ten_year = totalReturnWithDividends(candles, dividends, end_date.subtractYears(10), end_date),
    };
}

/// Morningstar-style trailing returns using month-end reference dates.
/// End date = last calendar day of prior month. Start date = that month-end minus N years.
/// Both dates snap backward to the last trading day on or before, matching
/// Morningstar's "last business day of the month" convention.
pub fn trailingReturnsMonthEnd(candles: []const Candle, today: Date) TrailingReturns {
    if (candles.len == 0) return .{};

    // End reference = last day of the prior month (snaps backward to last trading day)
    const month_end = today.lastDayOfPriorMonth();

    return .{
        .one_year = totalReturnFromAdjCloseBackward(candles, month_end.subtractYears(1), month_end),
        .three_year = totalReturnFromAdjCloseBackward(candles, month_end.subtractYears(3), month_end),
        .five_year = totalReturnFromAdjCloseBackward(candles, month_end.subtractYears(5), month_end),
        .ten_year = totalReturnFromAdjCloseBackward(candles, month_end.subtractYears(10), month_end),
    };
}

/// Same as trailingReturnsMonthEnd but with dividend reinvestment.
pub fn trailingReturnsMonthEndWithDividends(
    candles: []const Candle,
    dividends: []const Dividend,
    today: Date,
) TrailingReturns {
    if (candles.len == 0) return .{};

    const month_end = today.lastDayOfPriorMonth();

    return .{
        .one_year = totalReturnWithDividendsBackward(candles, dividends, month_end.subtractYears(1), month_end),
        .three_year = totalReturnWithDividendsBackward(candles, dividends, month_end.subtractYears(3), month_end),
        .five_year = totalReturnWithDividendsBackward(candles, dividends, month_end.subtractYears(5), month_end),
        .ten_year = totalReturnWithDividendsBackward(candles, dividends, month_end.subtractYears(10), month_end),
    };
}

const SearchDirection = enum { forward, backward };

/// Maximum calendar days a snapped candle can be from the target date.
/// Covers weekends + holidays (e.g., Christmas week). Beyond this, the data
/// is likely missing and the result would be misleading.
const max_snap_days: i32 = 10;

fn findNearestCandle(candles: []const Candle, target: Date, direction: SearchDirection) ?Candle {
    if (candles.len == 0) return null;

    // Binary search: lo = first index where candles[lo].date >= target
    var lo: usize = 0;
    var hi: usize = candles.len;
    while (lo < hi) {
        const mid = lo + (hi - lo) / 2;
        if (candles[mid].date.lessThan(target)) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }

    const candidate = switch (direction) {
        // First candle on or after target
        .forward => if (lo < candles.len) candles[lo] else return null,
        // Last candle on or before target
        .backward => if (lo < candles.len and candles[lo].date.eql(target))
            candles[lo]
        else if (lo > 0)
            candles[lo - 1]
        else
            return null,
    };

    // Reject if the snap distance exceeds tolerance
    const gap = candidate.date.days - target.days;
    if (gap > max_snap_days or gap < -max_snap_days) return null;

    return candidate;
}

/// Format a return value as a percentage string (e.g., "12.34%")
pub fn formatReturn(buf: []u8, value: f64) []const u8 {
    return std.fmt.bufPrint(buf, "{d:.2}%", .{value * 100.0}) catch "??%";
}

test "total return simple" {
    const candles = [_]Candle{
        .{ .date = Date.fromYmd(2024, 1, 2), .open = 100, .high = 101, .low = 99, .close = 100, .adj_close = 100, .volume = 1000 },
        .{ .date = Date.fromYmd(2024, 6, 28), .open = 110, .high = 111, .low = 109, .close = 110, .adj_close = 110, .volume = 1000 },
        .{ .date = Date.fromYmd(2024, 12, 31), .open = 120, .high = 121, .low = 119, .close = 120, .adj_close = 120, .volume = 1000 },
    };
    const result = totalReturnFromAdjClose(&candles, Date.fromYmd(2024, 1, 1), Date.fromYmd(2025, 1, 1));
    try std.testing.expect(result != null);
    // 120/100 - 1 = 0.20
    try std.testing.expectApproxEqAbs(@as(f64, 0.20), result.?.total_return, 0.001);
}

test "total return with dividends -- single dividend" {
    // Stock at $100, pays $2 dividend, price stays $100.
    // Without reinvestment: 0% return.
    // With reinvestment: $2/$100 = 0.02 extra shares -> 1.02 * $100 / $100 - 1 = 2%
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2024, 1, 2), 100),
        makeCandle(Date.fromYmd(2024, 3, 15), 100),
        makeCandle(Date.fromYmd(2024, 12, 31), 100),
    };
    const divs = [_]Dividend{
        .{ .ex_date = Date.fromYmd(2024, 3, 15), .amount = 2.0 },
    };
    const result = totalReturnWithDividends(&candles, &divs, Date.fromYmd(2024, 1, 1), Date.fromYmd(2025, 1, 1));
    try std.testing.expect(result != null);
    try std.testing.expectApproxEqAbs(@as(f64, 0.02), result.?.total_return, 0.0001);
}

test "total return with dividends -- quarterly dividends" {
    // Stock at $100 all year, pays $1 quarterly. Each $1 reinvested at $100 = 0.01 shares.
    // After Q1: 1.01 shares
    // After Q2: 1.01 + 1.01*1/100 = 1.01 * 1.01 = 1.0201
    // After Q3: 1.0201 * 1.01 = 1.030301
    // After Q4: 1.030301 * 1.01 = 1.04060401
    // Total return: 4.06%
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2024, 1, 2), 100),
        makeCandle(Date.fromYmd(2024, 3, 15), 100),
        makeCandle(Date.fromYmd(2024, 6, 14), 100),
        makeCandle(Date.fromYmd(2024, 9, 13), 100),
        makeCandle(Date.fromYmd(2024, 12, 13), 100),
        makeCandle(Date.fromYmd(2024, 12, 31), 100),
    };
    const divs = [_]Dividend{
        .{ .ex_date = Date.fromYmd(2024, 3, 15), .amount = 1.0 },
        .{ .ex_date = Date.fromYmd(2024, 6, 14), .amount = 1.0 },
        .{ .ex_date = Date.fromYmd(2024, 9, 13), .amount = 1.0 },
        .{ .ex_date = Date.fromYmd(2024, 12, 13), .amount = 1.0 },
    };
    const result = totalReturnWithDividends(&candles, &divs, Date.fromYmd(2024, 1, 1), Date.fromYmd(2025, 1, 1));
    try std.testing.expect(result != null);
    // (1.01)^4 - 1 = 0.04060401
    try std.testing.expectApproxEqAbs(@as(f64, 0.04060401), result.?.total_return, 0.0001);
}

test "total return with dividends -- price growth plus dividends" {
    // Start $100, end $120 (20% price return).
    // One $3 dividend at mid-year when price is $110.
    // Shares: 1 + 3/110 = 1.027273
    // Final value: 1.027273 * 120 = 123.2727
    // Total return: 123.2727 / 100 - 1 = 23.27%
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2024, 1, 2), 100),
        makeCandle(Date.fromYmd(2024, 6, 14), 110),
        makeCandle(Date.fromYmd(2024, 12, 31), 120),
    };
    const divs = [_]Dividend{
        .{ .ex_date = Date.fromYmd(2024, 6, 14), .amount = 3.0 },
    };
    const result = totalReturnWithDividends(&candles, &divs, Date.fromYmd(2024, 1, 1), Date.fromYmd(2025, 1, 1));
    try std.testing.expect(result != null);
    const expected = (1.0 + 3.0 / 110.0) * 120.0 / 100.0 - 1.0; // 0.23272727...
    try std.testing.expectApproxEqAbs(expected, result.?.total_return, 0.0001);
}

test "annualized return -- 3 year period" {
    // 3 years: $100 -> $150. Total return = 50%.
    // Annualized = (1.50)^(1/3) - 1 = 14.47%
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2021, 1, 4), 100),
        makeCandle(Date.fromYmd(2024, 1, 2), 150),
    };
    const result = totalReturnFromAdjClose(&candles, Date.fromYmd(2021, 1, 1), Date.fromYmd(2024, 1, 3));
    try std.testing.expect(result != null);
    try std.testing.expectApproxEqAbs(@as(f64, 0.50), result.?.total_return, 0.001);
    const ann = result.?.annualized_return.?;
    // (1.50)^(1/years) - 1, years ~ 3.0 (via 365.25)
    const years = Date.yearsBetween(Date.fromYmd(2021, 1, 4), Date.fromYmd(2024, 1, 2));
    const expected_ann = std.math.pow(f64, 1.50, 1.0 / years) - 1.0;
    try std.testing.expectApproxEqAbs(expected_ann, ann, 0.0001);
}

test "findNearestCandle -- exact match" {
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2024, 1, 2), 100),
        makeCandle(Date.fromYmd(2024, 1, 3), 101),
        makeCandle(Date.fromYmd(2024, 1, 4), 102),
    };
    // Forward exact
    const fwd = findNearestCandle(&candles, Date.fromYmd(2024, 1, 3), .forward).?;
    try std.testing.expect(fwd.date.eql(Date.fromYmd(2024, 1, 3)));
    // Backward exact
    const bwd = findNearestCandle(&candles, Date.fromYmd(2024, 1, 3), .backward).?;
    try std.testing.expect(bwd.date.eql(Date.fromYmd(2024, 1, 3)));
}

test "findNearestCandle -- weekend snap" {
    // Jan 4 2025 is Saturday, Jan 5 is Sunday
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2025, 1, 3), 100), // Friday
        makeCandle(Date.fromYmd(2025, 1, 6), 101), // Monday
    };
    // Forward from Saturday -> Monday
    const fwd = findNearestCandle(&candles, Date.fromYmd(2025, 1, 4), .forward).?;
    try std.testing.expect(fwd.date.eql(Date.fromYmd(2025, 1, 6)));
    // Backward from Saturday -> Friday
    const bwd = findNearestCandle(&candles, Date.fromYmd(2025, 1, 4), .backward).?;
    try std.testing.expect(bwd.date.eql(Date.fromYmd(2025, 1, 3)));
}

test "month-end trailing returns -- date windowing" {
    // Verify month-end logic uses correct reference dates.
    // "Today" = 2026-02-15, prior month end = 2026-01-31
    // 1yr window: 2025-01-31 to 2026-01-31
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2025, 1, 31), 100), // Jan 31 2025 is Friday
        makeCandle(Date.fromYmd(2025, 7, 1), 110),
        makeCandle(Date.fromYmd(2026, 1, 30), 120), // Jan 31 is Sat, trading day is 30th
        makeCandle(Date.fromYmd(2026, 2, 14), 125),
    };
    const today = Date.fromYmd(2026, 2, 15);
    const ret = trailingReturnsMonthEnd(&candles, today);
    // Month-end = Jan 31 2026. backward snap -> Jan 30.
    // Start = Jan 31 2025 (exact match, backward snap). End = Jan 30 2026.
    // Return = 120/100 - 1 = 20%
    try std.testing.expect(ret.one_year != null);
    try std.testing.expectApproxEqAbs(@as(f64, 0.20), ret.one_year.?.total_return, 0.001);
}

test "month-end trailing returns -- weekend start snaps backward" {
    // When the start month-end falls on a weekend, it should snap BACKWARD
    // to the last trading day (Friday), not forward to Monday.
    // This matches Morningstar's "last business day of the month" convention.
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2016, 1, 29), 100), // Friday (last biz day of Jan 2016)
        makeCandle(Date.fromYmd(2016, 2, 1), 95), // Monday (NOT what we want)
        makeCandle(Date.fromYmd(2026, 1, 30), 240), // End: Friday (last biz day of Jan 2026)
    };
    // Jan 31 2016 is Sunday. Backward snap -> Jan 29 (Friday).
    // Jan 31 2026 is Saturday. Backward snap -> Jan 30 (Friday).
    // Return = 240/100 - 1 = 140%
    const today = Date.fromYmd(2026, 2, 15);
    const ret = trailingReturnsMonthEnd(&candles, today);
    try std.testing.expect(ret.ten_year != null);
    try std.testing.expectApproxEqAbs(@as(f64, 1.40), ret.ten_year.?.total_return, 0.001);
    // Verify start date is Jan 29 (Friday), not Feb 1 (Monday)
    try std.testing.expect(ret.ten_year.?.from.eql(Date.fromYmd(2016, 1, 29)));
}

test "dividends outside window are excluded" {
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2024, 1, 2), 100),
        makeCandle(Date.fromYmd(2024, 6, 14), 100),
        makeCandle(Date.fromYmd(2024, 12, 31), 100),
    };
    const divs = [_]Dividend{
        .{ .ex_date = Date.fromYmd(2023, 12, 15), .amount = 5.0 }, // before window
        .{ .ex_date = Date.fromYmd(2024, 6, 14), .amount = 2.0 }, // inside
        .{ .ex_date = Date.fromYmd(2025, 3, 15), .amount = 5.0 }, // after window
    };
    const result = totalReturnWithDividends(&candles, &divs, Date.fromYmd(2024, 1, 1), Date.fromYmd(2025, 1, 1));
    try std.testing.expect(result != null);
    // Only the $2 mid-year dividend counts: 2/100 = 2%
    try std.testing.expectApproxEqAbs(@as(f64, 0.02), result.?.total_return, 0.0001);
}

test "zero price candle returns null" {
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2024, 1, 2), 0),
        makeCandle(Date.fromYmd(2024, 12, 31), 100),
    };
    const result = totalReturnFromAdjClose(&candles, Date.fromYmd(2024, 1, 1), Date.fromYmd(2025, 1, 1));
    try std.testing.expect(result == null);
}

test "empty candles returns null" {
    const candles = [_]Candle{};
    const result = totalReturnFromAdjClose(&candles, Date.fromYmd(2024, 1, 1), Date.fromYmd(2025, 1, 1));
    try std.testing.expect(result == null);
}

fn makeCandle(date: Date, price: f64) Candle {
    return .{ .date = date, .open = price, .high = price, .low = price, .close = price, .adj_close = price, .volume = 1000 };
}

// Morningstar reference data, captured 2026-02-24.
//
// AMZN Trailing Returns (as-of-date, from morningstar.com/stocks/xnas/amzn/trailing-returns):
//   Day end 2026-02-24:  1yr=-1.95%  3yr=30.66%  5yr=5.71%  10yr=22.37%
//   AMZN has no dividends, so price return = total return.
//
// VTI Trailing Returns (as-of-date, from morningstar.com/etfs/arcx/vti/trailing-returns):
//   Day end 2026-02-24:  1yr=16.62%  3yr=21.01%  5yr=12.03%  10yr=15.10% (price)
//
// VTI Performance (month-end, from morningstar.com/etfs/arcx/vti/performance):
//   Month-end Jan 31:    10yr total=15.10%  3yr total=20.20% (NAV ~20.24%)

test "as-of-date trailing returns -- AMZN vs Morningstar" {
    // Real AMZN split-adjusted closing prices from Twelve Data.
    // AMZN pays no dividends, so adj_close == close.
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2016, 2, 24), 27.702), // 10yr start
        makeCandle(Date.fromYmd(2021, 2, 24), 157.9765), // 5yr start
        makeCandle(Date.fromYmd(2023, 2, 24), 93.50), // 3yr start
        makeCandle(Date.fromYmd(2025, 2, 24), 212.71), // 1yr start
        makeCandle(Date.fromYmd(2026, 2, 24), 208.56), // end (latest close)
    };

    const ret = trailingReturns(&candles);

    // 1yr: 208.56 / 212.71 - 1 = -1.95%
    try std.testing.expect(ret.one_year != null);
    try std.testing.expectApproxEqAbs(@as(f64, -0.0195), ret.one_year.?.total_return, 0.001);

    // 3yr: annualized. Morningstar shows 30.66%.
    try std.testing.expect(ret.three_year != null);
    try std.testing.expectApproxEqAbs(@as(f64, 0.3066), ret.three_year.?.annualized_return.?, 0.002);

    // 5yr: annualized. Morningstar shows 5.71%.
    try std.testing.expect(ret.five_year != null);
    try std.testing.expectApproxEqAbs(@as(f64, 0.0571), ret.five_year.?.annualized_return.?, 0.002);

    // 10yr: annualized. Morningstar shows 22.37%.
    try std.testing.expect(ret.ten_year != null);
    try std.testing.expectApproxEqAbs(@as(f64, 0.2237), ret.ten_year.?.annualized_return.?, 0.002);
}

test "as-of-date vs month-end -- different results from same data" {
    // Demonstrates that as-of-date and month-end give different results
    // when the latest close differs significantly from the month-end close.
    //
    // "Today" = 2026-02-25, month-end = Jan 31 2026
    // As-of end = Feb 24 (latest candle), month-end = Jan 30 (snap from Jan 31 Sat)
    const candles = [_]Candle{
        makeCandle(Date.fromYmd(2025, 1, 31), 100), // month-end 1yr start (Friday)
        makeCandle(Date.fromYmd(2025, 2, 24), 100), // as-of 1yr start
        makeCandle(Date.fromYmd(2025, 7, 1), 110),
        makeCandle(Date.fromYmd(2026, 1, 30), 115), // month-end end (Friday, Jan 31 is Sat)
        makeCandle(Date.fromYmd(2026, 2, 24), 120), // as-of end (latest)
    };

    // As-of-date: end=Feb 24 ($120), start=Feb 24 prior year ($100) → 20%
    const asof = trailingReturns(&candles);
    try std.testing.expect(asof.one_year != null);
    try std.testing.expectApproxEqAbs(@as(f64, 0.20), asof.one_year.?.total_return, 0.001);

    // Month-end: end=Jan 30 ($115), start=Jan 31 ($100) → 15%
    const me = trailingReturnsMonthEnd(&candles, Date.fromYmd(2026, 2, 25));
    try std.testing.expect(me.one_year != null);
    try std.testing.expectApproxEqAbs(@as(f64, 0.15), me.one_year.?.total_return, 0.001);
}