const std = @import("std");
const Date = @import("date.zig").Date;

/// Type of holding in a portfolio lot.
pub const LotType = enum {
    stock, // stocks and ETFs (default)
    option, // option contracts
    cd, // certificates of deposit
    cash, // cash/money market
    illiquid, // illiquid assets (real estate, vehicles, etc.)
    watch, // watchlist item (no position, just track price)

    pub fn label(self: LotType) []const u8 {
        return switch (self) {
            .stock => "Stock",
            .option => "Option",
            .cd => "CD",
            .cash => "Cash",
            .illiquid => "Illiquid",
            .watch => "Watch",
        };
    }

    pub fn fromString(s: []const u8) LotType {
        if (std.mem.eql(u8, s, "option")) return .option;
        if (std.mem.eql(u8, s, "cd")) return .cd;
        if (std.mem.eql(u8, s, "cash")) return .cash;
        if (std.mem.eql(u8, s, "illiquid")) return .illiquid;
        if (std.mem.eql(u8, s, "watch")) return .watch;
        return .stock;
    }
};

/// A single lot in a portfolio -- one purchase/sale event.
/// Open lots have no close_date/close_price.
/// Closed lots have both.
pub const Lot = struct {
    symbol: []const u8,
    shares: f64,
    open_date: Date,
    open_price: f64,
    close_date: ?Date = null,
    close_price: ?f64 = null,
    /// Optional note/tag for the lot
    note: ?[]const u8 = null,
    /// Optional account identifier (e.g. "Roth IRA", "Brokerage")
    account: ?[]const u8 = null,
    /// Type of holding (stock, option, cd, cash)
    lot_type: LotType = .stock,
    /// Maturity date (for CDs)
    maturity_date: ?Date = null,
    /// Interest rate (for CDs, as percentage e.g. 3.8 = 3.8%)
    rate: ?f64 = null,
    /// Whether this lot is from dividend reinvestment (DRIP).
    /// DRIP lots are summarized as ST/LT groups instead of shown individually.
    drip: bool = false,
    /// Ticker alias for price fetching (e.g. CUSIP symbol with ticker::VTTHX).
    /// When set, this ticker is used for API calls instead of the symbol field.
    ticker: ?[]const u8 = null,
    /// Manual price override (e.g. for mutual funds not covered by data providers).
    /// Used as fallback when API price fetch fails.
    price: ?f64 = null,
    /// Date of the manual price (for display/staleness tracking).
    price_date: ?Date = null,

    /// The symbol to use for price fetching (ticker if set, else symbol).
    pub fn priceSymbol(self: Lot) []const u8 {
        return self.ticker orelse self.symbol;
    }

    pub fn isOpen(self: Lot) bool {
        return self.close_date == null;
    }

    pub fn costBasis(self: Lot) f64 {
        return self.shares * self.open_price;
    }

    pub fn marketValue(self: Lot, current_price: f64) f64 {
        return self.shares * current_price;
    }

    pub fn realizedPnl(self: Lot) ?f64 {
        const cp = self.close_price orelse return null;
        return self.shares * (cp - self.open_price);
    }

    pub fn unrealizedPnl(self: Lot, current_price: f64) f64 {
        return self.shares * (current_price - self.open_price);
    }

    pub fn returnPct(self: Lot, current_price: f64) f64 {
        if (self.open_price == 0) return 0;
        const price = if (self.close_price) |cp| cp else current_price;
        return (price / self.open_price) - 1.0;
    }
};

/// Aggregated position for a single symbol across multiple lots.
pub const Position = struct {
    symbol: []const u8,
    /// Total open shares
    shares: f64,
    /// Weighted average cost basis per share (open lots only)
    avg_cost: f64,
    /// Total cost basis of open lots
    total_cost: f64,
    /// Number of open lots
    open_lots: u32,
    /// Number of closed lots
    closed_lots: u32,
    /// Total realized P&L from closed lots
    realized_pnl: f64,
    /// Account name (shared across lots, or "Multiple" if mixed).
    account: []const u8 = "",
    /// Note from the first lot (e.g. "VANGUARD TARGET 2035").
    note: ?[]const u8 = null,
};

/// A portfolio is a collection of lots.
pub const Portfolio = struct {
    lots: []Lot,
    allocator: std.mem.Allocator,

    pub fn deinit(self: *Portfolio) void {
        for (self.lots) |lot| {
            self.allocator.free(lot.symbol);
            if (lot.note) |n| self.allocator.free(n);
            if (lot.account) |a| self.allocator.free(a);
            if (lot.ticker) |t| self.allocator.free(t);
        }
        self.allocator.free(self.lots);
    }

    /// Get all unique symbols in the portfolio (all types).
    pub fn symbols(self: Portfolio, allocator: std.mem.Allocator) ![][]const u8 {
        var seen = std.StringHashMap(void).init(allocator);
        defer seen.deinit();

        for (self.lots) |lot| {
            try seen.put(lot.symbol, {});
        }

        var result = std.ArrayList([]const u8).empty;
        errdefer result.deinit(allocator);

        var iter = seen.keyIterator();
        while (iter.next()) |key| {
            try result.append(allocator, key.*);
        }
        return result.toOwnedSlice(allocator);
    }

    /// Get unique symbols for stock/ETF lots only (skips options, CDs, cash).
    /// Returns the price symbol (ticker alias if set, otherwise raw symbol).
    pub fn stockSymbols(self: Portfolio, allocator: std.mem.Allocator) ![][]const u8 {
        var seen = std.StringHashMap(void).init(allocator);
        defer seen.deinit();

        for (self.lots) |lot| {
            if (lot.lot_type == .stock) {
                try seen.put(lot.priceSymbol(), {});
            }
        }

        var result = std.ArrayList([]const u8).empty;
        errdefer result.deinit(allocator);

        var iter = seen.keyIterator();
        while (iter.next()) |key| {
            try result.append(allocator, key.*);
        }
        return result.toOwnedSlice(allocator);
    }

    /// Get all lots for a given symbol.
    pub fn lotsForSymbol(self: Portfolio, allocator: std.mem.Allocator, symbol: []const u8) ![]Lot {
        var result = std.ArrayList(Lot).empty;
        errdefer result.deinit(allocator);

        for (self.lots) |lot| {
            if (std.mem.eql(u8, lot.symbol, symbol)) {
                try result.append(allocator, lot);
            }
        }
        return result.toOwnedSlice(allocator);
    }

    /// Get all lots of a given security type (allocated copy).
    pub fn lotsOfTypeAlloc(self: Portfolio, allocator: std.mem.Allocator, sec_type: LotType) ![]Lot {
        var result = std.ArrayList(Lot).empty;
        errdefer result.deinit(allocator);

        for (self.lots) |lot| {
            if (lot.lot_type == sec_type) {
                try result.append(allocator, lot);
            }
        }
        return result.toOwnedSlice(allocator);
    }

    /// Aggregate stock/ETF lots into positions by symbol (skips options, CDs, cash).
    /// Keys by priceSymbol() so CUSIP lots with ticker aliases aggregate under the ticker.
    pub fn positions(self: Portfolio, allocator: std.mem.Allocator) ![]Position {
        var map = std.StringHashMap(Position).init(allocator);
        defer map.deinit();

        for (self.lots) |lot| {
            if (lot.lot_type != .stock) continue;
            const sym = lot.priceSymbol();
            const entry = try map.getOrPut(sym);
            if (!entry.found_existing) {
                entry.value_ptr.* = .{
                    .symbol = sym,
                    .shares = 0,
                    .avg_cost = 0,
                    .total_cost = 0,
                    .open_lots = 0,
                    .closed_lots = 0,
                    .realized_pnl = 0,
                    .account = lot.account orelse "",
                    .note = lot.note,
                };
            } else {
                // Track account: if lots have different accounts, mark as "Multiple"
                const existing = entry.value_ptr.account;
                const new_acct = lot.account orelse "";
                if (existing.len > 0 and !std.mem.eql(u8, existing, "Multiple") and !std.mem.eql(u8, existing, new_acct)) {
                    entry.value_ptr.account = "Multiple";
                }
            }
            if (lot.isOpen()) {
                entry.value_ptr.shares += lot.shares;
                entry.value_ptr.total_cost += lot.costBasis();
                entry.value_ptr.open_lots += 1;
            } else {
                entry.value_ptr.closed_lots += 1;
                entry.value_ptr.realized_pnl += lot.realizedPnl() orelse 0;
            }
        }

        // Compute avg_cost
        var iter = map.valueIterator();
        while (iter.next()) |pos| {
            if (pos.shares > 0) {
                pos.avg_cost = pos.total_cost / pos.shares;
            }
        }

        var result = std.ArrayList(Position).empty;
        errdefer result.deinit(allocator);

        var viter = map.valueIterator();
        while (viter.next()) |pos| {
            try result.append(allocator, pos.*);
        }
        return result.toOwnedSlice(allocator);
    }

    /// Total cost basis of all open stock lots.
    pub fn totalCostBasis(self: Portfolio) f64 {
        var total: f64 = 0;
        for (self.lots) |lot| {
            if (lot.isOpen() and lot.lot_type == .stock) total += lot.costBasis();
        }
        return total;
    }

    /// Total realized P&L from all closed stock lots.
    pub fn totalRealizedPnl(self: Portfolio) f64 {
        var total: f64 = 0;
        for (self.lots) |lot| {
            if (lot.lot_type == .stock) {
                if (lot.realizedPnl()) |pnl| total += pnl;
            }
        }
        return total;
    }

    /// Total cash across all accounts.
    pub fn totalCash(self: Portfolio) f64 {
        var total: f64 = 0;
        for (self.lots) |lot| {
            if (lot.lot_type == .cash) total += lot.shares;
        }
        return total;
    }

    /// Total illiquid asset value across all accounts.
    pub fn totalIlliquid(self: Portfolio) f64 {
        var total: f64 = 0;
        for (self.lots) |lot| {
            if (lot.lot_type == .illiquid) total += lot.shares;
        }
        return total;
    }

    /// Total CD face value across all accounts.
    pub fn totalCdFaceValue(self: Portfolio) f64 {
        var total: f64 = 0;
        for (self.lots) |lot| {
            if (lot.lot_type == .cd) total += lot.shares;
        }
        return total;
    }

    /// Total option cost basis (absolute value of shares * open_price).
    pub fn totalOptionCost(self: Portfolio) f64 {
        var total: f64 = 0;
        for (self.lots) |lot| {
            if (lot.lot_type == .option) {
                // shares can be negative (short), open_price is per-contract cost
                total += @abs(lot.shares) * lot.open_price;
            }
        }
        return total;
    }

    /// Check if portfolio has any lots of a given type.
    pub fn hasType(self: Portfolio, sec_type: LotType) bool {
        for (self.lots) |lot| {
            if (lot.lot_type == sec_type) return true;
        }
        return false;
    }

    /// Get watchlist symbols (from watch lots in the portfolio).
    pub fn watchSymbols(self: Portfolio, allocator: std.mem.Allocator) ![][]const u8 {
        var result = std.ArrayList([]const u8).empty;
        errdefer result.deinit(allocator);

        for (self.lots) |lot| {
            if (lot.lot_type == .watch) {
                try result.append(allocator, lot.symbol);
            }
        }
        return result.toOwnedSlice(allocator);
    }
};

test "lot basics" {
    const lot = Lot{
        .symbol = "AAPL",
        .shares = 10,
        .open_date = Date.fromYmd(2024, 1, 15),
        .open_price = 150.0,
    };
    try std.testing.expect(lot.isOpen());
    try std.testing.expectApproxEqAbs(@as(f64, 1500.0), lot.costBasis(), 0.01);
    try std.testing.expectApproxEqAbs(@as(f64, 2000.0), lot.marketValue(200.0), 0.01);
    try std.testing.expectApproxEqAbs(@as(f64, 500.0), lot.unrealizedPnl(200.0), 0.01);
    try std.testing.expect(lot.realizedPnl() == null);
}

test "closed lot" {
    const lot = Lot{
        .symbol = "AAPL",
        .shares = 10,
        .open_date = Date.fromYmd(2024, 1, 15),
        .open_price = 150.0,
        .close_date = Date.fromYmd(2024, 6, 15),
        .close_price = 200.0,
    };
    try std.testing.expect(!lot.isOpen());
    try std.testing.expectApproxEqAbs(@as(f64, 500.0), lot.realizedPnl().?, 0.01);
    try std.testing.expectApproxEqAbs(@as(f64, 0.3333), lot.returnPct(0), 0.001);
}

test "portfolio positions" {
    const allocator = std.testing.allocator;

    var lots = [_]Lot{
        .{ .symbol = "AAPL", .shares = 10, .open_date = Date.fromYmd(2024, 1, 1), .open_price = 150.0 },
        .{ .symbol = "AAPL", .shares = 5, .open_date = Date.fromYmd(2024, 3, 1), .open_price = 160.0 },
        .{ .symbol = "VTI", .shares = 100, .open_date = Date.fromYmd(2024, 1, 1), .open_price = 220.0 },
        .{ .symbol = "AAPL", .shares = 3, .open_date = Date.fromYmd(2023, 6, 1), .open_price = 130.0, .close_date = Date.fromYmd(2024, 2, 1), .close_price = 155.0 },
    };

    var portfolio = Portfolio{ .lots = &lots, .allocator = allocator };
    // Don't call deinit since these are stack-allocated test strings

    const pos = try portfolio.positions(allocator);
    defer allocator.free(pos);

    try std.testing.expectEqual(@as(usize, 2), pos.len);

    // Find AAPL position
    var aapl: ?Position = null;
    for (pos) |p| {
        if (std.mem.eql(u8, p.symbol, "AAPL")) aapl = p;
    }
    try std.testing.expect(aapl != null);
    try std.testing.expectApproxEqAbs(@as(f64, 15.0), aapl.?.shares, 0.01);
    try std.testing.expectEqual(@as(u32, 2), aapl.?.open_lots);
    try std.testing.expectEqual(@as(u32, 1), aapl.?.closed_lots);
    try std.testing.expectApproxEqAbs(@as(f64, 75.0), aapl.?.realized_pnl, 0.01); // 3 * (155-130)
}