zfin/src/commands/enrich.zig

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

const OverviewMeta = struct {
    sector: []const u8,
    geo: []const u8,
    asset_class: []const u8,
};

/// Derive sector, geo, and asset_class from an Alpha Vantage company overview.
fn deriveMetadata(overview: zfin.CompanyOverview, sector_buf: []u8) OverviewMeta {
    const sector_raw = overview.sector orelse "Unknown";
    const sector_str = cli.fmt.toTitleCase(sector_buf, sector_raw);
    const country_str = overview.country orelse "US";
    const geo_str = if (std.mem.eql(u8, country_str, "USA")) "US" else country_str;

    const asset_class_str = blk: {
        if (overview.asset_type) |at| {
            if (std.mem.eql(u8, at, "ETF")) break :blk "ETF";
            if (std.mem.eql(u8, at, "Mutual Fund")) break :blk "Mutual Fund";
        }
        if (overview.market_cap) |mc_str| {
            const mc = std.fmt.parseInt(u64, mc_str, 10) catch 0;
            if (mc >= 10_000_000_000) break :blk "US Large Cap";
            if (mc >= 2_000_000_000) break :blk "US Mid Cap";
            break :blk "US Small Cap";
        }
        break :blk "US Large Cap";
    };

    return .{ .sector = sector_str, .geo = geo_str, .asset_class = asset_class_str };
}

/// CLI `enrich` command: bootstrap a metadata.srf file from Alpha Vantage OVERVIEW data.
/// Reads the portfolio, extracts stock symbols, fetches sector/industry/country for each,
/// and outputs a metadata SRF file to stdout.
/// If the argument looks like a symbol (no path separators, no .srf extension), enrich just that symbol.
pub fn run(allocator: std.mem.Allocator, svc: *zfin.DataService, arg: []const u8, out: *std.Io.Writer) !void {
    // Determine if arg is a symbol or a file path
    const is_file = std.mem.endsWith(u8, arg, ".srf") or
        std.mem.indexOfScalar(u8, arg, '/') != null or
        std.mem.indexOfScalar(u8, arg, '.') != null;

    if (!is_file) {
        // Single symbol mode: enrich one symbol, output appendable SRF (no header)
        try enrichSymbol(allocator, svc, arg, out);
        return;
    }

    // Portfolio file mode: enrich all symbols
    try enrichPortfolio(allocator, svc, arg, out);
}

/// Enrich a single symbol and output appendable SRF lines to stdout.
fn enrichSymbol(allocator: std.mem.Allocator, svc: *zfin.DataService, sym: []const u8, out: *std.Io.Writer) !void {
    {
        var msg_buf: [128]u8 = undefined;
        const msg = std.fmt.bufPrint(&msg_buf, "  Fetching {s}...\n", .{sym}) catch "  ...\n";
        try cli.stderrPrint(msg);
    }

    const overview = svc.getCompanyOverview(sym) catch |err| {
        if (err == zfin.DataError.NoApiKey) {
            try cli.stderrPrint("Error: ALPHAVANTAGE_API_KEY not set. Add it to .env\n");
            return;
        }
        try cli.stderrPrint("Error: Failed to fetch data for symbol\n");
        try out.print("# {s} -- fetch failed\n", .{sym});
        try out.print("# symbol::{s},sector::TODO,geo::TODO,asset_class::TODO\n", .{sym});
        return;
    };
    defer {
        if (overview.name) |n| allocator.free(n);
        if (overview.sector) |s| allocator.free(s);
        if (overview.industry) |ind| allocator.free(ind);
        if (overview.country) |c| allocator.free(c);
        if (overview.market_cap) |mc| allocator.free(mc);
        if (overview.asset_type) |at| allocator.free(at);
    }

    var sector_buf: [64]u8 = undefined;
    const meta = deriveMetadata(overview, &sector_buf);

    if (overview.name) |name| {
        try out.print("# {s}\n", .{name});
    }
    try out.print("symbol::{s},sector::{s},geo::{s},asset_class::{s}\n", .{
        sym, meta.sector, meta.geo, meta.asset_class,
    });
}

/// Enrich all symbols from a portfolio file.
fn enrichPortfolio(allocator: std.mem.Allocator, svc: *zfin.DataService, file_path: []const u8, out: *std.Io.Writer) !void {

    // Load portfolio
    const file_data = std.fs.cwd().readFileAlloc(allocator, file_path, 10 * 1024 * 1024) catch {
        try cli.stderrPrint("Error: Cannot read portfolio file\n");
        return;
    };
    defer allocator.free(file_data);

    var portfolio = zfin.cache.deserializePortfolio(allocator, file_data) catch {
        try cli.stderrPrint("Error: Cannot parse portfolio file\n");
        return;
    };
    defer portfolio.deinit();

    // Get unique stock symbols (using display-oriented names)
    const positions = try portfolio.positions(allocator);
    defer allocator.free(positions);

    // Get unique price symbols (raw API symbols)
    const syms = try portfolio.stockSymbols(allocator);
    defer allocator.free(syms);

    try out.print("#!srfv1\n", .{});
    try out.print("# Portfolio classification metadata\n", .{});
    try out.print("# Generated from Alpha Vantage OVERVIEW data\n", .{});
    try out.print("# Edit as needed: sector, geo, asset_class, pct:num:N\n", .{});
    try out.print("#\n", .{});
    try out.print("# For ETFs/funds with multi-class exposure, add multiple lines\n", .{});
    try out.print("# with pct:num: values that sum to ~100\n\n", .{});

    var success: usize = 0;
    var skipped: usize = 0;
    var failed: usize = 0;

    for (syms, 0..) |sym, i| {
        // Skip CUSIPs and known non-stock symbols
        if (isCusipLike(sym)) {
            // Find the display name for this CUSIP
            const display: []const u8 = sym;
            var note: ?[]const u8 = null;
            for (positions) |pos| {
                if (std.mem.eql(u8, pos.symbol, sym)) {
                    if (pos.note) |n| {
                        note = n;
                    }
                    break;
                }
            }
            try out.print("# CUSIP {s}", .{sym});
            if (note) |n| try out.print(" ({s})", .{n});
            try out.print(" -- fill in manually\n", .{});
            try out.print("# symbol::{s},asset_class::TODO,geo::TODO\n\n", .{display});
            skipped += 1;
            continue;
        }

        // Progress to stderr
        {
            var msg_buf: [128]u8 = undefined;
            const msg = std.fmt.bufPrint(&msg_buf, "  [{d}/{d}] {s}...\n", .{ i + 1, syms.len, sym }) catch "  ...\n";
            try cli.stderrPrint(msg);
        }

        const overview = svc.getCompanyOverview(sym) catch {
            try out.print("# {s} -- fetch failed\n", .{sym});
            try out.print("# symbol::{s},sector::TODO,geo::TODO,asset_class::TODO\n\n", .{sym});
            failed += 1;
            continue;
        };
        // Free allocated strings from overview when done
        defer {
            if (overview.name) |n| allocator.free(n);
            if (overview.sector) |s| allocator.free(s);
            if (overview.industry) |ind| allocator.free(ind);
            if (overview.country) |c| allocator.free(c);
            if (overview.market_cap) |mc| allocator.free(mc);
            if (overview.asset_type) |at| allocator.free(at);
        }

        var sector_buf: [64]u8 = undefined;
        const meta = deriveMetadata(overview, &sector_buf);

        // Comment with the name for readability
        if (overview.name) |name| {
            try out.print("# {s}\n", .{name});
        }
        try out.print("symbol::{s},sector::{s},geo::{s},asset_class::{s}\n\n", .{
            sym, meta.sector, meta.geo, meta.asset_class,
        });
        success += 1;
    }

    // Summary comment
    try out.print("# ---\n", .{});
    try out.print("# Enriched {d} symbols ({d} success, {d} skipped, {d} failed)\n", .{
        syms.len, success, skipped, failed,
    });
    try out.print("# Review and edit this file, then save as metadata.srf\n", .{});
}