#!/usr/bin/env python3
"""
Backfill sector/industry/exchange/country for all active tickers in universe_ticker_registry
that currently have sector=NULL.

Usage:
    python scripts/backfill_registry_sector.py [--batch 50] [--dry-run]

Expected runtime: ~30 min for 9376 tickers at semaphore=5.
"""

import asyncio
import argparse
import logging
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker

from app.core.config import settings
from app.models.universe_snapshot import UniverseTickerRegistry
from app.services import company_metadata_service as cms

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)


async def backfill(batch_size: int, dry_run: bool) -> None:
    engine = create_async_engine(settings.DATABASE_URL, echo=False)
    Session = async_sessionmaker(engine, expire_on_commit=False)

    async with Session() as db:
        result = await db.execute(
            select(UniverseTickerRegistry.ticker)
            .where(
                UniverseTickerRegistry.is_active == True,
                UniverseTickerRegistry.sector == None,
            )
            .order_by(UniverseTickerRegistry.ticker)
        )
        tickers = [row[0] for row in result.fetchall()]

    logger.info("Found %d tickers with NULL sector", len(tickers))
    if dry_run:
        logger.info("Dry-run mode — no changes will be written.")
        return

    success = 0
    failed = 0

    for i in range(0, len(tickers), batch_size):
        batch = tickers[i:i + batch_size]
        logger.info("Batch %d/%d — processing %d tickers", i // batch_size + 1,
                    (len(tickers) + batch_size - 1) // batch_size, len(batch))

        async def _enrich(ticker: str) -> bool:
            async with Session() as db:
                try:
                    meta = await cms.get_metadata(db, ticker)
                    sector = meta.get("sector")
                    logger.info("  %s → sector=%s", ticker, sector or "NULL")
                    return bool(sector)
                except ValueError:
                    logger.debug("  %s → invalid ticker, skipping", ticker)
                    return False
                except Exception as e:
                    logger.warning("  %s → error: %s", ticker, e)
                    return False

        results = await asyncio.gather(*[_enrich(t) for t in batch])
        success += sum(results)
        failed += sum(1 for r in results if not r)

        # Brief pause between batches to be kind to yfinance
        await asyncio.sleep(1)

    logger.info("Done. success=%d  failed/skipped=%d  total=%d", success, failed, len(tickers))
    await engine.dispose()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch", type=int, default=50)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    asyncio.run(backfill(args.batch, args.dry_run))


if __name__ == "__main__":
    main()