# ============================================================
# Project: EDGAR Research Database
# Description: SEC EDGAR Research Infrastructure
#
# Authors:
# - Manel Kammoun — manel.kammoun@uqo.ca
# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca
# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
#
# Institution:
# Université du Québec en Outaouais (UQO)
#
# License:
# Research / Academic Use
# ============================================================

from __future__ import annotations

import io
import json
from datetime import datetime, timezone
from pathlib import Path

import httpx
import pandas as pd

from src.utils.constants import DATA_DIRS
from src.utils.logging import get_logger

logger = get_logger("acquisition.sp500")

WIKIPEDIA_SP500_URL = "https://en.wikipedia.org/wiki/List_of_S%26P_500_companies"

FALLBACK_TICKERS = [
    {"ticker": "AAPL", "company_name": "Apple Inc.", "cik": "", "sector": "Information Technology", "sub_industry": "Technology Hardware", "date_added": ""},
    {"ticker": "MSFT", "company_name": "Microsoft Corporation", "cik": "", "sector": "Information Technology", "sub_industry": "Systems Software", "date_added": ""},
    {"ticker": "GOOGL", "company_name": "Alphabet Inc.", "cik": "", "sector": "Communication Services", "sub_industry": "Interactive Media", "date_added": ""},
    {"ticker": "AMZN", "company_name": "Amazon.com Inc.", "cik": "", "sector": "Consumer Discretionary", "sub_industry": "Internet Retail", "date_added": ""},
    {"ticker": "META", "company_name": "Meta Platforms Inc.", "cik": "", "sector": "Communication Services", "sub_industry": "Interactive Media", "date_added": ""},
    {"ticker": "NVDA", "company_name": "NVIDIA Corporation", "cik": "", "sector": "Information Technology", "sub_industry": "Semiconductors", "date_added": ""},
    {"ticker": "TSLA", "company_name": "Tesla Inc.", "cik": "", "sector": "Consumer Discretionary", "sub_industry": "Automobile Manufacturers", "date_added": ""},
    {"ticker": "JPM", "company_name": "JPMorgan Chase & Co.", "cik": "", "sector": "Financials", "sub_industry": "Diversified Banks", "date_added": ""},
    {"ticker": "JNJ", "company_name": "Johnson & Johnson", "cik": "", "sector": "Health Care", "sub_industry": "Pharmaceuticals", "date_added": ""},
    {"ticker": "V", "company_name": "Visa Inc.", "cik": "", "sector": "Financials", "sub_industry": "Transaction Processing", "date_added": ""},
    {"ticker": "PG", "company_name": "Procter & Gamble Co.", "cik": "", "sector": "Consumer Staples", "sub_industry": "Personal Products", "date_added": ""},
    {"ticker": "UNH", "company_name": "UnitedHealth Group Inc.", "cik": "", "sector": "Health Care", "sub_industry": "Managed Health Care", "date_added": ""},
    {"ticker": "HD", "company_name": "The Home Depot Inc.", "cik": "", "sector": "Consumer Discretionary", "sub_industry": "Home Improvement Retail", "date_added": ""},
    {"ticker": "MA", "company_name": "Mastercard Inc.", "cik": "", "sector": "Financials", "sub_industry": "Transaction Processing", "date_added": ""},
    {"ticker": "DIS", "company_name": "The Walt Disney Company", "cik": "", "sector": "Communication Services", "sub_industry": "Movies & Entertainment", "date_added": ""},
    {"ticker": "BAC", "company_name": "Bank of America Corp.", "cik": "", "sector": "Financials", "sub_industry": "Diversified Banks", "date_added": ""},
    {"ticker": "XOM", "company_name": "Exxon Mobil Corporation", "cik": "", "sector": "Energy", "sub_industry": "Integrated Oil & Gas", "date_added": ""},
    {"ticker": "PFE", "company_name": "Pfizer Inc.", "cik": "", "sector": "Health Care", "sub_industry": "Pharmaceuticals", "date_added": ""},
    {"ticker": "CSCO", "company_name": "Cisco Systems Inc.", "cik": "", "sector": "Information Technology", "sub_industry": "Communications Equipment", "date_added": ""},
    {"ticker": "VZ", "company_name": "Verizon Communications Inc.", "cik": "", "sector": "Communication Services", "sub_industry": "Integrated Telecom", "date_added": ""},
]

COLUMN_MAP = {
    "Symbol": "ticker",
    "Security": "company_name",
    "CIK": "cik",
    "GICS Sector": "sector",
    "GICS Sub-Industry": "sub_industry",
    "Date added": "date_added",
}


class SP500Universe:

    def __init__(self, cache_dir: Path | None = None) -> None:
        self._cache_dir = cache_dir or DATA_DIRS["metadata"]
        self._cache_path = self._cache_dir / "sp500_constituents.json"

    def get_constituents(self) -> list[dict]:
        try:
            resp = httpx.get(
                WIKIPEDIA_SP500_URL,
                headers={"User-Agent": "Mozilla/5.0 (EDGAR-Research-DB)"},
                follow_redirects=True,
                timeout=30,
            )
            resp.raise_for_status()
            tables = pd.read_html(io.StringIO(resp.text))
            df = tables[0]
            df = df.rename(columns=COLUMN_MAP)

            keep_cols = list(COLUMN_MAP.values())
            available = [c for c in keep_cols if c in df.columns]
            df = df[available]

            df["cik"] = df["cik"].astype(str).str.strip()
            df["ticker"] = df["ticker"].astype(str).str.strip()
            df["date_added"] = df["date_added"].astype(str).str.strip()

            constituents = df.to_dict(orient="records")
            logger.info("Fetched %d S&P 500 constituents from Wikipedia", len(constituents))

        except Exception:
            logger.warning("Wikipedia scrape failed, using fallback ticker list", exc_info=True)
            constituents = [dict(entry) for entry in FALLBACK_TICKERS]

        self._save_cache(constituents)
        return constituents

    def load_cached(self) -> list[dict]:
        if not self._cache_path.exists():
            logger.info("No cached constituents found, fetching fresh data")
            return self.get_constituents()

        with open(self._cache_path) as f:
            data = json.load(f)

        logger.info(
            "Loaded %d cached constituents (fetched %s)",
            len(data["constituents"]),
            data.get("fetched_at", "unknown"),
        )
        return data["constituents"]

    def get_tickers(self) -> list[str]:
        constituents = self.load_cached()
        return [c["ticker"] for c in constituents]

    def _save_cache(self, constituents: list[dict]) -> None:
        self._cache_dir.mkdir(parents=True, exist_ok=True)
        payload = {
            "fetched_at": datetime.now(timezone.utc).isoformat(),
            "count": len(constituents),
            "constituents": constituents,
        }
        with open(self._cache_path, "w") as f:
            json.dump(payload, f, indent=2, default=str)
        logger.info("Cached %d constituents to %s", len(constituents), self._cache_path)