#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import yaml
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]

def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--snapshot", default="GatherTradingData.json")
    ap.add_argument("--contract", default="spec/15_account_snapshot_contract.yaml")
    args = ap.parse_args()

    snapshot_path = ROOT / args.snapshot
    contract_path = ROOT / args.contract

    if not snapshot_path.exists():
        print(f"Snapshot file not found: {snapshot_path}")
        return 1

    data_payload = json.loads(snapshot_path.read_text(encoding="utf-8"))
    data = data_payload.get("data", {})
    settings = data.get("settings", {})
    
    # 1. Total Asset & Prev Market Regime
    total_asset = settings.get("total_asset_krw", 0)
    prev_regime = settings.get("prev_market_regime", "NORMAL")

    # Define cash floor ratio based on regime
    # normal: 10%, overheated_or_event_week: 15%, risk_off: 25% (or min_cash_ratio 7%/10%/15%)
    # Let's assume standard target is 10% for normal, 15% for overheated, 25% for risk_off
    # But settings could provide weekly_target_cash_pct
    target_ratio = settings.get("weekly_target_cash_pct", 10.0) / 100.0
    if not settings.get("weekly_target_cash_pct"):
        if "risk_off" in prev_regime.lower():
            target_ratio = 0.15 # min 15%
        elif "overheated" in prev_regime.lower() or "event" in prev_regime.lower():
            target_ratio = 0.10 # min 10%
        else:
            target_ratio = 0.07 # min 7%
            
    target_cash_floor = total_asset * target_ratio

    # 2. Gather Cash from snapshot
    snap = data.get("account_snapshot", [])
    
    immediate_cash_map = {}
    settlement_cash_d2_map = {}
    
    # Track accounts by type
    general_accounts = set()
    restricted_accounts = set()

    for item in snap:
        acc = item.get("account")
        acc_type = item.get("account_type", "")
        
        # Handle broken characters for general account type
        is_general = False
        if acc_type:
            acc_type_str = str(acc_type)
            if "일반" in acc_type_str or "Ϲ" in acc_type_str:
                is_general = True
            elif "isa" in acc_type_str.lower() or "연금" in acc_type_str or "pension" in acc_type_str.lower():
                is_general = False
            else:
                is_general = True # Default to general if unknown

        if is_general:
            if acc:
                general_accounts.add(acc)
        else:
            if acc:
                restricted_accounts.add(acc)

        # Retrieve cash values if present
        imm_cash = item.get("immediate_cash")
        if imm_cash is not None:
            immediate_cash_map[acc] = max(immediate_cash_map.get(acc, 0.0), float(imm_cash))
        d2_cash = item.get("settlement_cash_d2")
        if d2_cash is not None:
            settlement_cash_d2_map[acc] = max(settlement_cash_d2_map.get(acc, 0.0), float(d2_cash))

    # Backup from settings for general account D+2 cash if empty
    settings_d2 = settings.get("settlement_cash_d2_krw")
    if not settlement_cash_d2_map and settings_d2 is not None:
        # Assign to a stub general account if no accounts exist
        acc_stub = list(general_accounts)[0] if general_accounts else "general_stub"
        general_accounts.add(acc_stub)
        settlement_cash_d2_map[acc_stub] = float(settings_d2)

    # 3. Sum up cash by account priority rules
    general_immediate = sum(immediate_cash_map.get(acc, 0.0) for acc in general_accounts)
    general_d2 = sum(settlement_cash_d2_map.get(acc, 0.0) for acc in general_accounts)
    
    restricted_cash = sum(immediate_cash_map.get(acc, 0.0) + settlement_cash_d2_map.get(acc, 0.0) for acc in restricted_accounts)

    # Cross account cash leak detection
    # If restricted cash is added to general cash or general accounts have restricted tags
    cross_account_cash_leak_count = 0
    overlap = general_accounts.intersection(restricted_accounts)
    if overlap:
        cross_account_cash_leak_count += len(overlap)

    # D+2 Cash Defense Rule Applied
    d2_cash_defense_rule_applied = True
    eligible_cash = general_immediate + general_d2
    
    cash_shortfall = max(0.0, target_cash_floor - eligible_cash)

    gate = "PASS" if cross_account_cash_leak_count == 0 else "FAIL"

    result = {
        "formula_id": "CASH_LEDGER_V2",
        "total_asset_krw": total_asset,
        "target_cash_floor_krw": target_cash_floor,
        "immediate_cash": general_immediate,
        "settlement_cash_d2": general_d2,
        "restricted_cash": restricted_cash,
        "cross_account_cash_leak_count": cross_account_cash_leak_count,
        "d2_cash_defense_rule_applied": d2_cash_defense_rule_applied,
        "cash_shortfall": cash_shortfall,
        "gate": gate
    }

    out_path = ROOT / "Temp" / "cash_ledger_v2.json"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")

    print(json.dumps(result, ensure_ascii=True, indent=2))
    
    # Log information separately
    print(f"Eligible Cash (General Immediate + D+2): {eligible_cash:,.0f} KRW")
    print(f"Restricted Cash (ISA + Pension): {restricted_cash:,.0f} KRW")
    print(f"Cash Shortfall against Target Floor ({target_cash_floor:,.0f} KRW): {cash_shortfall:,.0f} KRW")

    return 0 if gate == "PASS" else 1

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