#!/usr/bin/env python3
"""Fetch constituent companies for a GICS sub-industry.

Usage:
    fetch_sub_industry_universe.py --code 45301020           # Semiconductors
    fetch_sub_industry_universe.py --name "Application Software"
    fetch_sub_industry_universe.py --sector Technology       # All sub-industries
    fetch_sub_industry_universe.py --code 45301020 --min-cap 500 --output ./reports/screening/universe.json

Identifies all publicly traded companies in a GICS Level 4 sub-industry by:
1. Mapping sub-industry to sector ETF holdings (yfinance)
2. Filtering by GICS sector/industry classification
3. Cross-referencing with thematic ETF holdings
4. Applying minimum market cap filter

Output: JSON with company list (ticker, name, market_cap, sector, industry).
"""

import argparse
import json
import os
import sys
from datetime import datetime, timezone

try:
    import yfinance as yf
except ImportError:
    sys.stderr.write("Error: yfinance required. Run: pip install yfinance\n")
    sys.exit(1)


# GICS Level 4 Sub-Industry → ETF constituents source mapping
# Maps sub-industry codes to ETFs whose holdings can be used to find constituents
GICS_SUB_INDUSTRY_ETFS = {
    # Energy (10)
    "10101010": {"name": "Oil & Gas Drilling", "etfs": ["XLE", "OIH"]},
    "10101020": {"name": "Oil & Gas Equipment & Services", "etfs": ["OIH", "XLE"]},
    "10102010": {"name": "Integrated Oil & Gas", "etfs": ["XLE"]},
    "10102020": {"name": "Oil & Gas Exploration & Production", "etfs": ["XOP", "XLE"]},
    "10102030": {"name": "Oil & Gas Refining & Marketing", "etfs": ["XLE", "CRAK"]},
    "10102040": {"name": "Oil & Gas Storage & Transportation", "etfs": ["AMLP", "XLE"]},
    "10102050": {"name": "Coal & Consumable Fuels", "etfs": ["XLE"]},
    # Materials (15)
    "15101010": {"name": "Commodity Chemicals", "etfs": ["XLB"]},
    "15101020": {"name": "Diversified Chemicals", "etfs": ["XLB"]},
    "15101030": {
        "name": "Fertilizers & Agricultural Chemicals",
        "etfs": ["MOO", "XLB"],
    },
    "15101040": {"name": "Industrial Gases", "etfs": ["XLB"]},
    "15101050": {"name": "Specialty Chemicals", "etfs": ["XLB"]},
    "15102010": {"name": "Construction Materials", "etfs": ["XLB", "XHB"]},
    "15104020": {"name": "Diversified Metals & Mining", "etfs": ["XME", "XLB"]},
    "15104025": {"name": "Copper", "etfs": ["COPX", "XME"]},
    "15104030": {"name": "Gold", "etfs": ["GDX", "GDXJ"]},
    "15104050": {"name": "Steel", "etfs": ["SLX", "XME"]},
    # Industrials (20)
    "20101010": {"name": "Aerospace & Defense", "etfs": ["ITA", "XLI"]},
    "20102010": {"name": "Building Products", "etfs": ["XLI", "XHB"]},
    "20103010": {"name": "Construction & Engineering", "etfs": ["XLI"]},
    "20106010": {"name": "Construction Machinery & Heavy Equipment", "etfs": ["XLI"]},
    "20106020": {"name": "Industrial Machinery & Components", "etfs": ["XLI"]},
    "20301010": {"name": "Air Freight & Logistics", "etfs": ["IYT", "XLI"]},
    "20302010": {"name": "Passenger Airlines", "etfs": ["JETS"]},
    "20304010": {"name": "Rail Transportation", "etfs": ["IYT"]},
    "20304020": {"name": "Cargo Ground Transportation", "etfs": ["IYT"]},
    # Consumer Discretionary (25)
    "25102010": {"name": "Automobile Manufacturers", "etfs": ["CARZ", "XLY"]},
    "25201030": {"name": "Homebuilding", "etfs": ["ITB", "XHB"]},
    "25203010": {"name": "Apparel, Accessories & Luxury Goods", "etfs": ["XRT", "XLY"]},
    "25301010": {"name": "Casinos & Gaming", "etfs": ["BJK", "XLY"]},
    "25301020": {"name": "Hotels, Resorts & Cruise Lines", "etfs": ["PEJ", "XLY"]},
    "25301040": {"name": "Restaurants", "etfs": ["PEJ", "XLY"]},
    "25502020": {"name": "Broadline Retail", "etfs": ["XRT", "XLY"]},
    "25503030": {"name": "Home Improvement Retail", "etfs": ["XHB", "XRT"]},
    # Consumer Staples (30)
    "30201030": {
        "name": "Soft Drinks & Non-alcoholic Beverages",
        "etfs": ["PBJ", "XLP"],
    },
    "30202030": {"name": "Packaged Foods & Meats", "etfs": ["PBJ", "XLP"]},
    "30301010": {"name": "Household Products", "etfs": ["XLP"]},
    # Health Care (35)
    "35101010": {"name": "Health Care Equipment", "etfs": ["IHI", "XLV"]},
    "35102030": {"name": "Managed Health Care", "etfs": ["IHF", "XLV"]},
    "35201010": {"name": "Biotechnology", "etfs": ["XBI", "IBB"]},
    "35202010": {"name": "Pharmaceuticals", "etfs": ["XLV"]},
    "35203010": {"name": "Life Sciences Tools & Services", "etfs": ["XLV"]},
    # Financials (40)
    "40101010": {"name": "Diversified Banks", "etfs": ["XLF"]},
    "40101015": {"name": "Regional Banks", "etfs": ["KRE"]},
    "40201060": {
        "name": "Transaction & Payment Processing Services",
        "etfs": ["IPAY", "XLF"],
    },
    "40202010": {"name": "Consumer Finance", "etfs": ["XLF"]},
    "40203010": {"name": "Asset Management & Custody Banks", "etfs": ["XLF"]},
    "40203020": {"name": "Investment Banking & Brokerage", "etfs": ["XLF"]},
    "40203040": {"name": "Financial Exchanges & Data", "etfs": ["XLF"]},
    "40301010": {"name": "Insurance Brokers", "etfs": ["KIE", "XLF"]},
    "40301040": {"name": "Property & Casualty Insurance", "etfs": ["KIE", "XLF"]},
    # Information Technology (45)
    "45102010": {"name": "IT Consulting & Other Services", "etfs": ["IGV", "XLK"]},
    "45102030": {
        "name": "Data Processing & Outsourced Services",
        "etfs": ["IPAY", "XLK"],
    },
    "45103010": {"name": "Application Software", "etfs": ["IGV", "XLK"]},
    "45103020": {"name": "Systems Software", "etfs": ["IGV", "XLK", "CIBR"]},
    "45201020": {"name": "Communications Equipment", "etfs": ["XLK"]},
    "45202010": {"name": "Technology Hardware, Storage & Peripherals", "etfs": ["XLK"]},
    "45203010": {"name": "Electronic Equipment & Instruments", "etfs": ["XLK"]},
    "45203015": {"name": "Electronic Components", "etfs": ["XLK"]},
    "45301010": {
        "name": "Semiconductor Materials & Equipment",
        "etfs": ["SMH", "SOXX"],
    },
    "45301020": {"name": "Semiconductors", "etfs": ["SMH", "SOXX"]},
    # Communication Services (50)
    "50101020": {"name": "Integrated Telecommunication Services", "etfs": ["XLC"]},
    "50201010": {"name": "Advertising", "etfs": ["XLC"]},
    "50202010": {"name": "Movies & Entertainment", "etfs": ["XLC"]},
    "50202020": {"name": "Interactive Home Entertainment", "etfs": ["HERO", "XLC"]},
    "50203010": {"name": "Interactive Media & Services", "etfs": ["XLC"]},
    # Utilities (55)
    "55101010": {"name": "Electric Utilities", "etfs": ["XLU"]},
    "55103010": {"name": "Multi-Utilities", "etfs": ["XLU"]},
    "55105010": {
        "name": "Independent Power Producers & Energy Traders",
        "etfs": ["XLU"],
    },
    "55105020": {"name": "Renewable Electricity", "etfs": ["ICLN", "QCLN"]},
    # Real Estate (60)
    "60102510": {"name": "Industrial REITs", "etfs": ["VNQ", "XLRE"]},
    "60104010": {"name": "Office REITs", "etfs": ["VNQ", "XLRE"]},
    "60106010": {"name": "Multi-Family Residential REITs", "etfs": ["VNQ", "XLRE"]},
    "60107010": {"name": "Retail REITs", "etfs": ["VNQ", "XLRE"]},
    "60108020": {"name": "Data Center REITs", "etfs": ["VNQ", "XLRE"]},
    "60108030": {"name": "Self-Storage REITs", "etfs": ["VNQ", "XLRE"]},
    "60108040": {"name": "Telecom Tower REITs", "etfs": ["VNQ", "XLRE"]},
}

# Name → code lookup (lowercase for matching)
NAME_TO_CODE = {v["name"].lower(): k for k, v in GICS_SUB_INDUSTRY_ETFS.items()}

# Sector → sub-industry codes
SECTOR_TO_CODES = {
    "Energy": [c for c in GICS_SUB_INDUSTRY_ETFS if c.startswith("10")],
    "Materials": [c for c in GICS_SUB_INDUSTRY_ETFS if c.startswith("15")],
    "Industrials": [c for c in GICS_SUB_INDUSTRY_ETFS if c.startswith("20")],
    "Consumer Discretionary": [c for c in GICS_SUB_INDUSTRY_ETFS if c.startswith("25")],
    "Consumer Staples": [c for c in GICS_SUB_INDUSTRY_ETFS if c.startswith("30")],
    "Healthcare": [c for c in GICS_SUB_INDUSTRY_ETFS if c.startswith("35")],
    "Financials": [c for c in GICS_SUB_INDUSTRY_ETFS if c.startswith("40")],
    "Technology": [c for c in GICS_SUB_INDUSTRY_ETFS if c.startswith("45")],
    "Communication Services": [c for c in GICS_SUB_INDUSTRY_ETFS if c.startswith("50")],
    "Utilities": [c for c in GICS_SUB_INDUSTRY_ETFS if c.startswith("55")],
    "Real Estate": [c for c in GICS_SUB_INDUSTRY_ETFS if c.startswith("60")],
}


def get_etf_holdings(etf_ticker: str) -> list[dict]:
    """Fetch ETF top holdings using yfinance."""
    try:
        etf = yf.Ticker(etf_ticker)
        info = etf.info or {}

        holdings = []
        # Try to get holdings from fund data
        if hasattr(etf, "funds_data"):
            try:
                fund_holdings = etf.funds_data.top_holdings
                if fund_holdings is not None and not fund_holdings.empty:
                    for ticker_sym in fund_holdings.index:
                        holdings.append({"ticker": str(ticker_sym)})
            except Exception:
                pass

        return holdings
    except Exception:
        return []


def get_company_info(ticker: str) -> dict | None:
    """Get basic company info for a ticker."""
    try:
        stock = yf.Ticker(ticker)
        info = stock.info or {}
        if not info.get("shortName") and not info.get("longName"):
            return None

        return {
            "ticker": ticker,
            "name": info.get("shortName") or info.get("longName", ticker),
            "market_cap": info.get("marketCap"),
            "sector": info.get("sector", ""),
            "industry": info.get("industry", ""),
            "country": info.get("country", ""),
        }
    except Exception:
        return None


def fetch_universe(gics_code: str, min_cap_millions: float = 500) -> dict:
    """Fetch all companies in a GICS sub-industry."""
    sub_info = GICS_SUB_INDUSTRY_ETFS.get(gics_code)
    if not sub_info:
        return {"error": f"Unknown GICS code: {gics_code}"}

    sub_name = sub_info["name"]
    etfs = sub_info["etfs"]

    # Collect candidate tickers from ETF holdings
    all_tickers = set()
    for etf_ticker in etfs:
        holdings = get_etf_holdings(etf_ticker)
        for h in holdings:
            t = h.get("ticker", "")
            if t and len(t) <= 5 and t.isalpha():
                all_tickers.add(t.upper())

    # Get company info and filter by market cap
    companies = []
    filtered_out = []
    min_cap = min_cap_millions * 1_000_000

    for ticker in sorted(all_tickers):
        info = get_company_info(ticker)
        if not info:
            continue

        mc = info.get("market_cap")
        if mc and mc < min_cap:
            filtered_out.append(
                {
                    "ticker": ticker,
                    "reason": f"Market cap ${mc/1e9:.1f}B < ${min_cap_millions}M minimum",
                }
            )
            continue

        companies.append(info)

    # Sort by market cap descending
    companies.sort(key=lambda x: x.get("market_cap") or 0, reverse=True)

    return {
        "gics_code": gics_code,
        "sub_industry": sub_name,
        "etf_sources": etfs,
        "universe_size": len(companies),
        "filtered_out_count": len(filtered_out),
        "min_market_cap_millions": min_cap_millions,
        "companies": companies,
        "filtered_out": filtered_out[:10],
        "computed_at": datetime.now(timezone.utc).isoformat(),
        "methodology": (
            f"Universe constructed from ETF holdings ({', '.join(etfs)}), "
            f"filtered by market cap >= ${min_cap_millions}M. "
            "Cross-reference with GICS classification for accuracy."
        ),
    }


def main():
    parser = argparse.ArgumentParser(
        description="Fetch constituent companies for a GICS Level 4 sub-industry"
    )
    parser.add_argument(
        "--code", help="8-digit GICS sub-industry code (e.g., 45301020)"
    )
    parser.add_argument("--name", help="Sub-industry name (e.g., 'Semiconductors')")
    parser.add_argument(
        "--sector", help="Fetch all sub-industries for a sector (e.g., 'Technology')"
    )
    parser.add_argument(
        "--min-cap",
        type=float,
        default=500,
        help="Minimum market cap in millions (default: 500)",
    )
    parser.add_argument("--output", help="Output file path (default: stdout)")
    parser.add_argument(
        "--list", action="store_true", help="List all known sub-industry codes"
    )
    args = parser.parse_args()

    if args.list:
        print("Available GICS Sub-Industry Codes:")
        print("-" * 60)
        for code, info in sorted(GICS_SUB_INDUSTRY_ETFS.items()):
            print(f"  {code}  {info['name']}")
        sys.exit(0)

    if args.sector:
        codes = SECTOR_TO_CODES.get(args.sector)
        if not codes:
            sys.stderr.write(
                f"Error: Unknown sector '{args.sector}'. "
                f"Available: {', '.join(SECTOR_TO_CODES.keys())}\n"
            )
            sys.exit(1)
        output = {"sector": args.sector, "sub_industries": {}}
        for code in codes:
            result = fetch_universe(code, args.min_cap)
            output["sub_industries"][code] = result
    elif args.code:
        output = fetch_universe(args.code, args.min_cap)
    elif args.name:
        code = NAME_TO_CODE.get(args.name.lower())
        if not code:
            sys.stderr.write(
                f"Error: Unknown sub-industry name '{args.name}'. "
                f"Use --list to see available names.\n"
            )
            sys.exit(1)
        output = fetch_universe(code, args.min_cap)
    else:
        sys.stderr.write("Error: Provide --code, --name, --sector, or --list\n")
        sys.exit(1)

    result_json = json.dumps(output, indent=2)
    if args.output:
        os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
        with open(args.output, "w") as f:
            f.write(result_json)
    else:
        print(result_json)
    sys.exit(0)


if __name__ == "__main__":
    main()