#!/usr/bin/env python3
from __future__ import annotations
"""
Collect real financial data for a US publicly traded company.
Output: structured JSON consumable by downstream financial skills.

Usage:
    python collect_data.py TICKER [--years 5] [--output path/to/output.json]

Examples:
    python collect_data.py META
    python collect_data.py AAPL --years 3 --output /tmp/aapl_data.json
    python collect_data.py NVDA --years 5

Default mode uses Python standard library plus public Yahoo Finance web endpoints.
Optional enhancement mode uses yfinance + pandas when available.
"""
# /// script
# requires-python = ">=3.11"
# dependencies = []
# ///

import argparse
import json
import importlib
import sys
import time
from datetime import date
from pathlib import Path

import yahoo_web_api

pd = None
yf = None

# Field name aliases — yfinance row indices vary across versions
FIELD_ALIASES = {
    "revenue": ["Total Revenue", "Revenue", "Operating Revenue"],
    "ebit": ["Operating Income", "EBIT"],
    "ebitda": ["EBITDA", "Normalized EBITDA"],
    "tax_expense": ["Tax Provision", "Income Tax Expense"],
    "net_income": ["Net Income", "Net Income Common Stockholders"],
    "capex": ["Capital Expenditure", "Capital Expenditures"],
    "operating_cash_flow": ["Operating Cash Flow", "Cash Flow From Continuing Operating Activities"],
    "depreciation_amortization": ["Depreciation And Amortization", "Depreciation Amortization Depletion"],
    "free_cash_flow": ["Free Cash Flow"],
    "change_in_nwc": ["Change In Working Capital", "Changes In Working Capital"],
    "sbc": ["Stock Based Compensation"],
    "total_debt": ["Total Debt"],
    "long_term_debt": ["Long Term Debt"],
    "short_term_debt": ["Short Long Term Debt", "Current Debt"],
    "cash": ["Cash And Cash Equivalents"],
    "short_investments": ["Other Short Term Investments", "Short Term Investments"],
    "current_assets": ["Current Assets"],
    "current_liabilities": ["Current Liabilities"],
    "total_assets": ["Total Assets"],
    "total_equity": ["Stockholders Equity", "Total Equity Gross Minority Interest"],
}


def load_optional_yfinance() -> bool:
    global pd, yf
    if pd is not None and yf is not None:
        return True
    try:
        pd = importlib.import_module("pandas")
        yf = importlib.import_module("yfinance")
        return True
    except ImportError:
        return False


def safe_get(df, field_key: str, col) -> float | None:
    """Safely extract a value from a DataFrame using alias chain."""
    if pd is None:
        raise RuntimeError("pandas is required for yfinance mode")
    aliases = FIELD_ALIASES.get(field_key, [field_key])
    for alias in aliases:
        if alias in df.index:
            val = df.loc[alias, col]
            if pd.notna(val):
                return float(val)
    return None


def get_year_col(df: pd.DataFrame, year: int):
    """Find the column matching a target year in a yfinance DataFrame."""
    matches = [c for c in df.columns if c.year == year]
    return matches[0] if matches else None


def collect_market_data(ticker_obj) -> dict:
    """Collect real-time market data."""
    info = ticker_obj.info
    price = info.get("currentPrice") or info.get("regularMarketPrice")
    shares_raw = info.get("sharesOutstanding")
    shares = shares_raw / 1e6 if shares_raw else None
    mcap_raw = info.get("marketCap")
    mcap = mcap_raw / 1e6 if mcap_raw else None
    beta = info.get("beta")

    return {
        "current_price": price,
        "shares_outstanding_millions": round(shares, 2) if shares else None,
        "market_cap_millions": round(mcap, 2) if mcap else None,
        "beta_5y_monthly": round(beta, 3) if beta else None,
    }


def collect_risk_free_rate() -> float | None:
    """Get 10Y Treasury yield as risk-free rate proxy."""
    if yf is None:
        return None
    try:
        tnx = yf.Ticker("^TNX")
        hist = tnx.history(period="5d")
        if hist.empty:
            return None
        yield_pct = hist["Close"].iloc[-1]
        return round(float(yield_pct) / 100, 4)  # Convert percentage to decimal
    except Exception:
        return None


def collect_income_statement(ticker_obj, years: list[int]) -> dict:
    """Collect income statement data for specified years."""
    try:
        financials = ticker_obj.financials
    except Exception:
        return {str(y): {"_source": "yfinance error"} for y in years}

    if financials is None or financials.empty:
        return {str(y): {"_source": "yfinance returned empty"} for y in years}

    result = {}
    nan_years = []
    for year in years:
        col = get_year_col(financials, year)
        if col is None:
            result[str(year)] = {"_source": f"yfinance has no column for {year}"}
            nan_years.append(year)
            continue

        revenue = safe_get(financials, "revenue", col)
        ebit = safe_get(financials, "ebit", col)

        if revenue is None and ebit is None:
            nan_years.append(year)
            result[str(year)] = {
                "revenue": None, "ebit": None, "ebitda": None,
                "tax_expense": None, "net_income": None,
                "_source": f"yfinance returned NaN for {year} — supplement from 10-K filing",
            }
        else:
            result[str(year)] = {
                "revenue": revenue,
                "ebit": ebit,
                "ebitda": safe_get(financials, "ebitda", col),
                "tax_expense": safe_get(financials, "tax_expense", col),
                "net_income": safe_get(financials, "net_income", col),
                "_source": "yfinance",
            }

    return result, nan_years


def collect_cash_flow(ticker_obj, years: list[int]) -> dict:
    """Collect cash flow data."""
    try:
        cf = ticker_obj.cashflow
    except Exception:
        return {str(y): {"_source": "yfinance error"} for y in years}, []

    if cf is None or cf.empty:
        return {str(y): {"_source": "yfinance returned empty"} for y in years}, years

    result = {}
    nan_years = []
    for year in years:
        col = get_year_col(cf, year)
        if col is None:
            result[str(year)] = {"_source": f"yfinance has no column for {year}"}
            nan_years.append(year)
            continue

        ocf = safe_get(cf, "operating_cash_flow", col)
        capex = safe_get(cf, "capex", col)

        if ocf is None and capex is None:
            nan_years.append(year)
            result[str(year)] = {
                "operating_cash_flow": None, "capex": None,
                "depreciation_amortization": None, "free_cash_flow": None,
                "change_in_nwc": None, "sbc": None,
                "_source": f"yfinance returned NaN for {year}",
            }
        else:
            result[str(year)] = {
                "operating_cash_flow": ocf,
                "capex": capex,  # Negative = outflow (preserved)
                "depreciation_amortization": safe_get(cf, "depreciation_amortization", col),
                "free_cash_flow": safe_get(cf, "free_cash_flow", col),
                "change_in_nwc": safe_get(cf, "change_in_nwc", col),
                "sbc": safe_get(cf, "sbc", col),
                "_source": "yfinance",
            }

    return result, nan_years


def collect_balance_sheet(ticker_obj, latest_year: int) -> dict:
    """Collect balance sheet data for the latest year."""
    try:
        bs = ticker_obj.balance_sheet
    except Exception:
        return {str(latest_year): {"_source": "yfinance error"}}

    if bs is None or bs.empty:
        return {str(latest_year): {"_source": "yfinance returned empty"}}

    col = get_year_col(bs, latest_year)
    if col is None:
        return {str(latest_year): {"_source": f"yfinance has no column for {latest_year}"}}

    total_debt = safe_get(bs, "total_debt", col)
    if total_debt is None:
        lt = safe_get(bs, "long_term_debt", col) or 0
        st = safe_get(bs, "short_term_debt", col) or 0
        total_debt = lt + st if (lt or st) else None

    cash = safe_get(bs, "cash", col)
    short_inv = safe_get(bs, "short_investments", col)
    cash_equiv = (cash or 0) + (short_inv or 0) if (cash is not None or short_inv is not None) else None
    net_debt = (total_debt - cash_equiv) if (total_debt is not None and cash_equiv is not None) else None

    return {
        str(latest_year): {
            "total_debt": total_debt,
            "cash_and_equivalents": cash_equiv,
            "net_debt": round(net_debt, 2) if net_debt is not None else None,
            "current_assets": safe_get(bs, "current_assets", col),
            "current_liabilities": safe_get(bs, "current_liabilities", col),
            "total_assets": safe_get(bs, "total_assets", col),
            "total_equity": safe_get(bs, "total_equity", col),
            "_source": "yfinance",
        }
    }


def collect_analyst_estimates(ticker_obj) -> dict:
    """Collect analyst consensus estimates."""
    result = {"revenue_next_fy": None, "revenue_fy_after": None, "eps_next_fy": None, "_source": "missing"}

    try:
        rev_est = ticker_obj.revenue_estimate
        if rev_est is not None and not rev_est.empty:
            if "0y" in rev_est.index:
                result["revenue_next_fy"] = float(rev_est.loc["0y", "avg"]) / 1e6 if pd.notna(rev_est.loc["0y", "avg"]) else None
            if "+1y" in rev_est.index:
                result["revenue_fy_after"] = float(rev_est.loc["+1y", "avg"]) / 1e6 if pd.notna(rev_est.loc["+1y", "avg"]) else None
            result["_source"] = "yfinance revenue_estimate"
    except Exception:
        pass

    try:
        eps = ticker_obj.eps_trend
        if eps is not None and not eps.empty:
            if "0y" in eps.index:
                val = eps.loc["0y", "current"] if "current" in eps.columns else eps.iloc[eps.index.get_loc("0y"), 0]
                result["eps_next_fy"] = float(val) if pd.notna(val) else None
    except Exception:
        pass

    return result


def infer_available_years_from_yfinance(ticker_obj, years_requested: int, fallback_latest_full_year: int) -> list[int]:
    try:
        financials = ticker_obj.financials
    except Exception:
        financials = None

    available_years = []
    if financials is not None and getattr(financials, "empty", True) is False:
        for col in financials.columns:
            year = getattr(col, "year", None)
            if year is not None:
                available_years.append(int(year))

    if available_years:
        unique_years = sorted(set(available_years))
        return unique_years[-years_requested:]

    return list(range(fallback_latest_full_year - years_requested + 1, fallback_latest_full_year + 1))


def resolve_engine(requested_engine: str, yfinance_available: bool) -> str:
    if requested_engine == "auto":
        if yfinance_available:
            return "yfinance"
        raise RuntimeError(
            "auto engine requires optional yfinance/pandas for full-quality collection. "
            "Install them or use the SEC helper script for low-dependency filing data."
        )
    if requested_engine == "yfinance" and not yfinance_available:
        raise RuntimeError("yfinance engine requested but yfinance/pandas are not installed")
    return requested_engine


def assemble_output(
    ticker_symbol: str,
    company_name: str,
    market_data: dict,
    income_data: dict,
    cf_data: dict,
    bs_data: dict,
    rfr: float | None,
    estimates: dict,
    all_nan_years: list[int],
    elapsed: float,
    target_years: list[int],
    engine_label: str,
):
    return {
        "ticker": ticker_symbol,
        "company_name": company_name,
        "data_date": date.today().isoformat(),
        "currency": "USD",
        "unit": "millions_usd",
        "data_sources": {
            "market_data": f"{engine_label} market data",
            "historical_financials": f"{engine_label} annual financials",
            "risk_free_rate": "yahoo finance ^TNX chart",
        },
        "market_data": market_data,
        "income_statement": income_data,
        "cash_flow": cf_data,
        "balance_sheet": bs_data,
        "wacc_inputs": {
            "risk_free_rate": rfr,
            "beta": market_data.get("beta_5y_monthly"),
            "credit_rating": None,
            "_source": f"{engine_label} + ^TNX",
        },
        "analyst_estimates": estimates,
        "metadata": {
            "_capex_convention": "negative = cash outflow (Yahoo convention)",
            "_fcf_note": "Yahoo FCF may differ from investment-bank FCF. Does NOT necessarily deduct SBC consistently.",
            "_nan_years": all_nan_years,
            "_collection_duration_seconds": elapsed,
            "_target_years": target_years,
            "_engine": engine_label,
        },
    }


def collect_with_yahoo_web(ticker_symbol: str, target_years: list[int]) -> tuple[dict, list[int]]:
    quote_summary = yahoo_web_api.fetch_quote_summary(ticker_symbol)
    parsed = yahoo_web_api.parse_quote_summary(ticker_symbol, quote_summary, target_years)
    market_data = parsed["market_data"]
    if market_data.get("current_price") is None:
        raise RuntimeError(f"{ticker_symbol} not found or no market data available from yahoo web")
    treasury_chart = yahoo_web_api.fetch_treasury_chart()
    rfr = yahoo_web_api.parse_treasury_chart(treasury_chart)
    return (
        assemble_output(
            ticker_symbol=ticker_symbol,
            company_name=parsed["company_name"],
            market_data=market_data,
            income_data=parsed["income_statement"],
            cf_data=parsed["cash_flow"],
            bs_data=parsed["balance_sheet"],
            rfr=rfr,
            estimates=parsed["analyst_estimates"],
            all_nan_years=parsed["metadata"].get("_nan_years", []),
            elapsed=0.0,
            target_years=target_years,
            engine_label="yahoo web quoteSummary",
        ),
        parsed["metadata"].get("_nan_years", []),
    )


def collect_with_yfinance(ticker_symbol: str, years_requested: int, fallback_latest_full_year: int) -> tuple[dict, list[int]]:
    if not load_optional_yfinance():
        raise RuntimeError("yfinance/pandas are not installed")

    ticker_obj = yf.Ticker(ticker_symbol)
    info = ticker_obj.info
    if not info or info.get("regularMarketPrice") is None:
        raise RuntimeError(f"{ticker_symbol} not found or no market data available")

    company_name = info.get("longName") or info.get("shortName") or ticker_symbol
    target_years = infer_available_years_from_yfinance(
        ticker_obj, years_requested=years_requested, fallback_latest_full_year=fallback_latest_full_year
    )
    market_data = collect_market_data(ticker_obj)
    rfr = collect_risk_free_rate()
    income_data, is_nan_years = collect_income_statement(ticker_obj, target_years)
    cf_data, cf_nan_years = collect_cash_flow(ticker_obj, target_years)
    bs_data = collect_balance_sheet(ticker_obj, max(target_years))
    estimates = collect_analyst_estimates(ticker_obj)
    all_nan_years = sorted(set(is_nan_years + cf_nan_years))
    return (
        assemble_output(
            ticker_symbol=ticker_symbol,
            company_name=company_name,
            market_data=market_data,
            income_data=income_data,
            cf_data=cf_data,
            bs_data=bs_data,
            rfr=rfr,
            estimates=estimates,
            all_nan_years=all_nan_years,
            elapsed=0.0,
            target_years=target_years,
            engine_label="yfinance",
        ),
        all_nan_years,
    )


def main():
    parser = argparse.ArgumentParser(description="Collect financial data for a US public company")
    parser.add_argument("ticker", help="Stock ticker symbol (e.g., META, AAPL)")
    parser.add_argument("--years", type=int, default=5, help="Number of historical years (default: 5)")
    parser.add_argument("--output", help="Output JSON file path (default: {TICKER}_financial_data.json)")
    parser.add_argument(
        "--engine",
        choices=["auto", "yahoo-web", "yfinance"],
        default="auto",
        help="Data collection engine. 'auto' defaults to standard-library yahoo-web mode.",
    )
    args = parser.parse_args()

    ticker_symbol = args.ticker.upper()
    current_year = date.today().year
    # Exclude current year if annual report not yet filed (typically before Q4 earnings)
    latest_full_year = current_year - 1
    target_years = list(range(latest_full_year - args.years + 1, latest_full_year + 1))

    output_path = args.output or f"{ticker_symbol}_financial_data.json"

    engine = resolve_engine(args.engine, yfinance_available=load_optional_yfinance())
    print(f"Collecting data for {ticker_symbol} ({target_years[0]}-{target_years[-1]}) using {engine}...")
    start_time = time.time()
    try:
        if engine == "yfinance":
            output, all_nan_years = collect_with_yfinance(
                ticker_symbol, years_requested=args.years, fallback_latest_full_year=latest_full_year
            )
        else:
            output, all_nan_years = collect_with_yahoo_web(ticker_symbol, target_years)
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        sys.exit(1)

    elapsed = round(time.time() - start_time, 1)
    output["metadata"]["_collection_duration_seconds"] = elapsed

    # Write output
    Path(output_path).write_text(json.dumps(output, indent=2, ensure_ascii=False))
    print(f"\nSaved to {output_path} ({elapsed}s)")

    # Summary
    if all_nan_years:
        print(f"\nWARNING: Years {all_nan_years} have missing data (NaN from yfinance).")
        print("   Supplement from 10-K filings or SEC EDGAR before using in models.")

    # Quick sanity check
    market_data = output["market_data"]
    if market_data["current_price"] and market_data["shares_outstanding_millions"]:
        computed_mcap = market_data["current_price"] * market_data["shares_outstanding_millions"]
        reported_mcap = market_data["market_cap_millions"]
        if reported_mcap and abs(computed_mcap - reported_mcap) / reported_mcap > 0.05:
            print(f"WARNING: Market cap mismatch: Price x Shares={computed_mcap:.0f}M vs Reported={reported_mcap:.0f}M")


if __name__ == "__main__":
    main()
