#!/usr/bin/env python3
import os
import time
from dotenv import load_dotenv
import requests

load_dotenv()

API_KEY = os.getenv("MASSIVE_API_KEY", "").strip()

BASE_URL = "https://api.massive.com"
REQUEST_TIMEOUT = 30
PAGE_LIMIT = 1000

LOGO_DIR = "logos"

FMP_LOGO_URL = "https://financialmodelingprep.com/image-stock/{symbol}.png"
IEX_LOGO_URL = "https://storage.googleapis.com/iex/api/logos/{symbol}.png"


def ensure_dirs():
    os.makedirs(LOGO_DIR, exist_ok=True)


def make_session():
    session = requests.Session()
    session.headers.update({
        "Authorization": f"Bearer {API_KEY}",
        "Accept": "application/json",
    })
    return session


def sanitize_symbol(symbol: str) -> str:
    return "".join(c if c.isalnum() or c in "-_." else "_" for c in symbol)


def get_logo_local_path(symbol: str) -> str:
    return os.path.join(LOGO_DIR, f"{sanitize_symbol(symbol)}.png")


def iter_all_symbols(session):
    url = f"{BASE_URL}/v3/reference/tickers"
    params = {
        "market": "stocks",
        "active": "true",
        "limit": PAGE_LIMIT,
    }

    while True:
        response = session.get(url, params=params, timeout=REQUEST_TIMEOUT)
        response.raise_for_status()
        data = response.json()

        for row in data.get("results", []):
            yield row

        if not data.get("next_url"):
            break

        url = data["next_url"]
        params = None


def get_overview(session, symbol: str) -> dict:
    url = f"{BASE_URL}/v3/reference/tickers/{symbol}"
    response = session.get(url, timeout=REQUEST_TIMEOUT)
    response.raise_for_status()
    return response.json().get("results", {})


def extract_logo(overview: dict) -> str | None:
    branding = overview.get("branding") or {}
    return branding.get("icon_url") or branding.get("logo_url")


def download_logo(session, url: str, symbol: str) -> bool:
    try:
        response = session.get(url, timeout=REQUEST_TIMEOUT)
        if response.status_code != 200:
            return False

        path = get_logo_local_path(symbol)
        with open(path, "wb") as f:
            f.write(response.content)
        return True
    except Exception:
        return False


def download_logo_with_fallback(session, overview: dict, symbol: str) -> bool:
    logo_url = extract_logo(overview)
    if logo_url and download_logo(session, logo_url, symbol):
        return True

    if download_logo(session, FMP_LOGO_URL.format(symbol=symbol), symbol):
        return True

    return download_logo(session, IEX_LOGO_URL.format(symbol=symbol), symbol)


def main():
    ensure_dirs()
    session = make_session()

    count = 0
    downloaded = 0

    for item in iter_all_symbols(session):
        symbol = item.get("ticker")
        if not symbol:
            continue

        try:
            overview = get_overview(session, symbol)
        except Exception:
            overview = {}

        if download_logo_with_fallback(session, overview, symbol):
            downloaded += 1

        count += 1
        if count % 200 == 0:
            print(f"Processed {count} symbols | Downloaded {downloaded} logos")

        time.sleep(0.005)

    print(f"Done. Processed {count} symbols | Downloaded {downloaded} logos")


if __name__ == "__main__":
    main()