#!/usr/bin/env python3
"""Position size calculator for US stocks/ETFs in KRW accounts.

Optional live mode uses yfinance if installed. Otherwise provide --price-usd
and --usdkrw manually. Every output includes the FX source and timestamp so
trade sizing is never based on a stale hidden assumption.
"""

from __future__ import annotations

import argparse
import math
import json
import sys
from datetime import datetime


LEVERAGED_OR_INVERSE = {
    "SOXL", "SOXS", "TQQQ", "SQQQ", "UPRO", "SPXL", "SPXS",
    "QLD", "SSO", "PSQ", "SH", "TECL", "TECS", "FNGU", "FNGD",
}


def fetch_yfinance(ticker: str) -> tuple[float | None, float | None, str]:
    try:
        import yfinance as yf  # type: ignore
    except Exception:
        return None, None, "yfinance not installed"

    try:
        stock = yf.Ticker(ticker)
        info = getattr(stock, "fast_info", None)
        price = None
        if info:
            price = info.get("last_price") or info.get("regular_market_price")
        if not price:
            hist = stock.history(period="1d", interval="1m")
            if not hist.empty:
                price = float(hist["Close"].dropna().iloc[-1])

        fx_ticker = yf.Ticker("KRW=X")
        fx_info = getattr(fx_ticker, "fast_info", None)
        usdkrw = None
        if fx_info:
            usdkrw = fx_info.get("last_price") or fx_info.get("regular_market_price")
        if not usdkrw:
            fx_hist = fx_ticker.history(period="1d", interval="1m")
            if not fx_hist.empty:
                usdkrw = float(fx_hist["Close"].dropna().iloc[-1])

        return float(price) if price else None, float(usdkrw) if usdkrw else None, "yfinance: ticker + KRW=X"
    except Exception as exc:
        return None, None, f"yfinance error: {exc}"


def main() -> int:
    parser = argparse.ArgumentParser(description="Calculate US stock share count from KRW budget.")
    parser.add_argument("ticker", help="Ticker, e.g. SOXL, SOXS, AMD")
    parser.add_argument("--account", choices=["main", "sub"], required=True)
    parser.add_argument("--budget-krw", type=float, required=True)
    parser.add_argument("--price-usd", type=float)
    parser.add_argument("--usdkrw", type=float)
    parser.add_argument("--price-source", default="manual")
    parser.add_argument("--fx-source", default="manual")
    parser.add_argument("--fee-buffer-pct", type=float, default=0.15)
    parser.add_argument("--live", action="store_true", help="Try yfinance for price and USD/KRW")
    parser.add_argument("--json", action="store_true", help="Print machine-readable JSON")
    args = parser.parse_args()

    ticker = args.ticker.upper()
    if args.account == "sub" and ticker in LEVERAGED_OR_INVERSE:
        print(f"NOT ELIGIBLE: {ticker} is leveraged/inverse and sub account cannot trade it.")
        return 2

    price = args.price_usd
    usdkrw = args.usdkrw
    price_source = args.price_source
    fx_source = args.fx_source

    if args.live:
        live_price, live_fx, live_source = fetch_yfinance(ticker)
        price = live_price or price
        usdkrw = live_fx or usdkrw
        price_source = live_source if live_price else price_source
        fx_source = live_source if live_fx else fx_source

    if not price or not usdkrw:
        print("NEEDS LIVE PRICE/FX: provide --price-usd and --usdkrw, or install yfinance and use --live.")
        return 1

    timestamp = datetime.now().isoformat(timespec="seconds")
    fee_multiplier = 1 + (args.fee_buffer_pct / 100)
    per_share_krw = price * usdkrw * fee_multiplier
    shares = math.floor(args.budget_krw / per_share_krw)
    exposure = shares * price * usdkrw
    leftover = args.budget_krw - (shares * per_share_krw)

    result = {
        "timestamp": timestamp,
        "account": args.account,
        "ticker": ticker,
        "price_source": price_source,
        "price_timestamp": timestamp,
        "fx_source": fx_source,
        "fx_timestamp": timestamp,
        "price_usd": round(price, 4),
        "usdkrw": round(usdkrw, 4),
        "budget_krw": round(args.budget_krw),
        "fee_buffer_pct": round(args.fee_buffer_pct, 4),
        "calculated_shares": shares,
        "approx_exposure_krw": round(exposure),
        "fee_buffered_leftover_krw": round(leftover),
    }

    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return 0

    print(f"timestamp: {result['timestamp']}")
    print(f"account: {result['account']}")
    print(f"ticker: {result['ticker']}")
    print(f"price_source: {result['price_source']}")
    print(f"fx_source: {result['fx_source']}")
    print(f"price_usd: {result['price_usd']:,.4f}")
    print(f"usdkrw: {result['usdkrw']:,.4f}")
    print(f"budget_krw: {result['budget_krw']:,.0f}")
    print(f"fee_buffer_pct: {result['fee_buffer_pct']:.3f}")
    print(f"calculated_shares: {result['calculated_shares']:,}")
    print(f"approx_exposure_krw: {result['approx_exposure_krw']:,.0f}")
    print(f"fee_buffered_leftover_krw: {result['fee_buffered_leftover_krw']:,.0f}")
    return 0


if __name__ == "__main__":
    sys.exit(main())