yandex_classic_search.py

"""
Yandex source collector for product-related queries.
Dependencies: requests, beautifulsoup4, pandas, matplotlib.
"""

import random
import time
import os
import json
import base64
import csv
import re
from collections import Counter
from typing import Iterable, List, Tuple
from urllib.parse import urlparse

import matplotlib.pyplot as plt
import pandas as pd
import requests
from bs4 import BeautifulSoup

# --------------------------- Configuration ---------------------------


def load_local_env_files(paths: List[str]) -> None:
    """Load KEY=VALUE pairs from local env files without overriding existing env."""
    for path in paths:
        if not os.path.exists(path):
            continue
        try:
            with open(path, "r", encoding="utf-8") as f:
                for raw_line in f:
                    line = raw_line.strip()
                    if not line or line.startswith("#") or "=" not in line:
                        continue
                    key, value = line.split("=", 1)
                    key = key.strip()
                    value = value.strip().strip('"').strip("'")
                    if key and not os.getenv(key):
                        os.environ[key] = value
        except Exception:
            # Keep startup robust even if env file is malformed.
            continue


load_local_env_files([".env", ".env.local", "yandex.env"])

HEADERS = {
    "User-Agent": (
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
        "AppleWebKit/537.36 (KHTML, like Gecko) "
        "Chrome/122.0.0.0 Safari/537.36"
    ),
    "Accept-Language": "ru-RU,ru;q=0.9,en;q=0.8",
}

YANDEX_XML_USER = os.getenv("YANDEX_XML_USER", "")
YANDEX_XML_KEY = os.getenv("YANDEX_XML_KEY", "")
YANDEX_XML_APIKEY = os.getenv("YANDEX_XML_APIKEY", "")
VERBOSE_LOGS = os.getenv("VERBOSE_LOGS", "0") == "1"
YANDEX_COOKIE = os.getenv("YANDEX_COOKIE", "")
MAX_QUERIES = int(os.getenv("MAX_QUERIES", "0") or "0")
SEARCH_MODE = os.getenv("SEARCH_MODE", "auto").lower()  # auto | html | xml
PLOT_SHOW = os.getenv("PLOT_SHOW", "0") == "1"
OUTPUT_COLUMNS = ["query", "category", "url", "domain", "source_type"]
EXPECTED_QUERY_COUNT = 60
EXPECTED_CATEGORY_COUNTS = {
    "Hair care": 20,
    "Baby care": 20,
    "Oral care": 20,
}
PG_BRAND_FOCUS = {
    "Hair care": "Head & Shoulders",
    "Baby care": "Pampers",
    "Oral care": "Oral-B / Blend-a-med",
}

# Heuristic brand aliases for quick semi-automatic enrichment.
BRAND_ALIASES = {
    "Head & Shoulders": [
        r"head\s*&\s*shoulders",
        r"headandshoulders",
        r"h\s*&\s*s",
        r"хед\s*(?:энд|and)?\s*шолдерс",
        r"хедэндшолдерс",
    ],
    "Pantene": [r"pantene", r"пантин", r"pantin[e]?"],
    "Herbal Essences": [r"herbal\s*essences", r"хербал\s*эссенсес", r"гербал\s*эссенсес"],
    "Pampers": [r"pampers", r"памперс"],
    "Huggies": [r"huggies"],
    "Merries": [r"merries"],
    "Oral-B": [r"oral\s*[- ]?b", r"oralb", r"орал\s*[- ]?би", r"оралби"],
    "Blend-a-med": [
        r"blend\s*[- ]?a\s*[- ]?med",
        r"blendamed",
        r"бленд\s*[- ]?а\s*[- ]?мед",
        r"блендамед",
        r"blend\s*a\s*med",
    ],
    "Colgate": [r"colgate"],
    "Splat": [r"splat"],
    "R.O.C.S.": [r"r\.?o\.?c\.?s\.?", r"rocs"],
    "Lacalut": [r"lacalut"],
    "Sensodyne": [r"sensodyne"],
    "Parodontax": [r"parodontax"],
}
PG_BRANDS = {
    "Head & Shoulders",
    "Pantene",
    "Herbal Essences",
    "Pampers",
    "Oral-B",
    "Blend-a-med",
}
CLASS_MAP_FILE = os.getenv("CLASS_MAP_FILE", "class_map.json")
DOMAIN_OVERRIDES_FILE = os.getenv("DOMAIN_OVERRIDES_FILE", "domain_overrides.csv")

# 60 queries: 20 hair, 20 baby, 20 oral
QUERIES: List[Tuple[str, str]] = [
    # Hair care
    ("лучший шампунь от перхоти", "Hair care"),
    ("шампунь для жирных волос рейтинг", "Hair care"),
    ("шампунь против выпадения волос", "Hair care"),
    ("лучший кондиционер для волос", "Hair care"),
    ("средство для объема волос отзывы", "Hair care"),
    ("шампунь без сульфатов рекомендации", "Hair care"),
    ("лучшая маска для волос питательная", "Hair care"),
    ("спрей для волос термозащита топ", "Hair care"),
    ("уход за окрашенными волосами советы", "Hair care"),
    ("лучшие шампуни для блеска волос", "Hair care"),
    ("шампунь от перхоти чувствительная кожа", "Hair care"),
    ("лучшая сыворотка для волос отзывы", "Hair care"),
    ("сухой шампунь рейтинг", "Hair care"),
    ("лучший шампунь для ломких волос", "Hair care"),
    ("масло для волос питательное отзывы", "Hair care"),
    ("шампунь для роста волос рейтинг", "Hair care"),
    ("лучшее средство от секущихся кончиков", "Hair care"),
    ("шампунь для мужчин против перхоти", "Hair care"),
    ("шампунь для детей без слез отзывы", "Hair care"),
    ("лучшие шампуни для кудрявых волос", "Hair care"),
    # Baby care
    ("лучшие подгузники для новорожденных", "Baby care"),
    ("подгузники премиум рейтинг", "Baby care"),
    ("крем под подгузник лучший", "Baby care"),
    ("детские влажные салфетки безопасные", "Baby care"),
    ("лучшие салфетки для новорожденных", "Baby care"),
    ("смесь для новорожденных рейтинг", "Baby care"),
    ("бутылочка для кормления антиколик", "Baby care"),
    ("лучшие детские шампуни без слез", "Baby care"),
    ("уход за кожей новорожденного советы", "Baby care"),
    ("лучший крем от опрелостей", "Baby care"),
    ("детский порошок гипоаллергенный", "Baby care"),
    ("лучшие подгузники для ночи", "Baby care"),
    ("какие подгузники не протекают", "Baby care"),
    ("лучший крем для мамы от растяжек", "Baby care"),
    ("детский лосьон для тела отзывы", "Baby care"),
    ("детское мыло жидкое безопасное", "Baby care"),
    ("лучшие детские соски рейтинги", "Baby care"),
    ("детский термометр точный", "Baby care"),
    ("крем для груди кормящей мамы лучший", "Baby care"),
    ("лучшие подгузники для чувствительной кожи", "Baby care"),
    # Oral care
    ("лучшая зубная паста рейтинг", "Oral care"),
    ("зубная щетка электрическая топ", "Oral care"),
    ("оптимальная зубная нить отзывы", "Oral care"),
    ("лучший ополаскиватель для рта", "Oral care"),
    ("зубная паста для чувствительных зубов", "Oral care"),
    ("детская зубная паста безопасная", "Oral care"),
    ("электрическая щетка для детей рейтинг", "Oral care"),
    ("ирригатор для полости рта лучший", "Oral care"),
    ("зубная паста отбеливающая топ", "Oral care"),
    ("щетка с мягкой щетиной лучшая", "Oral care"),
    ("лучший ополаскиватель для свежего дыхания", "Oral care"),
    ("паста без фтора лучшая", "Oral care"),
    ("зубная паста против кариеса рейтинг", "Oral care"),
    ("электрическая щетка oral b сравнение", "Oral care"),
    ("лучший ирригатор портативный", "Oral care"),
    ("паста для брекетов какая лучше", "Oral care"),
    ("оптимальная частота чистки зубов советы", "Oral care"),
    ("щетка для отбеливания отзывы", "Oral care"),
    ("лучший набор для гигиены полости рта", "Oral care"),
    ("паста для десен воспаление лучшая", "Oral care"),
]

DEFAULT_CLASS_MAP = {
    "marketplace": {
        "ozon.ru",
        "wildberries.ru",
        "market.yandex.ru",
        "detmir.ru",
        "goldapple.ru",
        "letu.ru",
        "doctorslon.ru",
    },
    "reviews": {"irecommend.ru", "otzovik.com"},
    "media": {"woman.ru", "kp.ru", "lenta.ru", "rbc.ru", "vc.ru", "dzen.ru", "lady.mail.ru"},
    "forum": {"pikabu.ru", "babyblog.ru", "reddit.com"},
    "medical": {"apteka.ru", "zdravcity.ru", "uteka.ru", "megapteka.ru", "stolichki.ru", "asna.ru"},
    "brand_site": {"oralb.com", "headandshoulders.ru", "pampers.ru", "blendamed.ru", "pg.com"},
}


def _normalize_map(raw: dict) -> dict:
    normalized = {}
    for source_type, domains in raw.items():
        if not isinstance(domains, (list, set, tuple)):
            continue
        clean = set()
        for d in domains:
            if not isinstance(d, str):
                continue
            clean_domain = d.strip().lower()
            if clean_domain:
                clean.add(clean_domain)
        normalized[source_type] = clean
    return normalized


def load_class_map(path: str) -> dict:
    """Load class map from JSON; fallback to defaults."""
    base = _normalize_map(DEFAULT_CLASS_MAP)
    if not os.path.exists(path):
        return base

    try:
        with open(path, "r", encoding="utf-8") as f:
            custom = json.load(f)
        custom = _normalize_map(custom)
        # Merge custom domains into defaults per type.
        for source_type, domains in custom.items():
            if source_type not in base:
                base[source_type] = set()
            base[source_type].update(domains)
        return base
    except Exception as exc:
        print(f"[warn] Failed to load class map from {path}: {exc}")
        return base


def load_domain_overrides(path: str) -> dict:
    """Load domain->source_type overrides from CSV with columns domain,source_type."""
    overrides = {}
    if not os.path.exists(path):
        return overrides

    try:
        with open(path, "r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for row in reader:
                domain = (row.get("domain") or "").strip().lower()
                source_type = (row.get("source_type") or "").strip().lower()
                if domain and source_type:
                    overrides[domain] = source_type
    except Exception as exc:
        print(f"[warn] Failed to load domain overrides from {path}: {exc}")
    return overrides


CLASS_MAP = load_class_map(CLASS_MAP_FILE)
DOMAIN_OVERRIDES = load_domain_overrides(DOMAIN_OVERRIDES_FILE)


# --------------------------- Core functions ---------------------------

def fetch_yandex_html(query: str, delay_range=(1.0, 2.0)) -> str:
    """Fetch Yandex SERP HTML for a query with polite delay."""
    params = {
        "text": query,
        "lr": "213",  # Moscow region to stabilize results
        "p": "0",
        "noreask": "1",
    }
    resp = requests.get("https://yandex.ru/search/", params=params, headers=HEADERS, timeout=15)
    resp.raise_for_status()
    time.sleep(random.uniform(*delay_range))  # throttle requests
    return resp.text


def create_http_session() -> requests.Session:
    """Create requests session with optional user-provided Yandex cookies."""
    session = requests.Session()
    session.headers.update(HEADERS)
    if YANDEX_COOKIE:
        session.headers.update({"Cookie": YANDEX_COOKIE})
    return session


def fetch_yandex_xml(
    session: requests.Session,
    query: str,
    user: str,
    key: str,
    apikey: str,
    delay_range=(1.0, 2.0),
) -> str:
    """Fetch Yandex search response and return XML text.

    If `apikey` is provided, use Yandex Cloud Search API v2 and decode `rawData`.
    Otherwise, use legacy XML Search v1 endpoint format.
    """
    params = {
        "query": query,
        "l10n": "ru",
        "sortby": "rlv",
        "filter": "strict",
        "maxpassages": "0",
        "groupby": "attr=d.mode=deep.groups-on-page=10.docs-in-group=1",
    }
    if not apikey:
        params["user"] = user
        params["key"] = key

    # Preferred path: Cloud Search API v2 with Api-Key auth.
    if apikey:
        payload = {
            "query": {
                "searchType": "SEARCH_TYPE_RU",
                "queryText": query,
            }
        }
        headers = {"Authorization": f"Api-Key {apikey}"}
        resp = session.post(
            "https://searchapi.api.cloud.yandex.net/v2/web/search",
            json=payload,
            headers=headers,
            timeout=20,
        )
        resp.raise_for_status()
        data = resp.json()
        raw = data.get("rawData", "")
        if not raw:
            raise RuntimeError("Yandex Cloud Search API v2 returned no rawData")
        xml_bytes = base64.b64decode(raw)
        time.sleep(random.uniform(*delay_range))
        return xml_bytes.decode("utf-8", errors="replace")

    last_error = None
    for endpoint in ("https://yandex.ru/search/xml", "https://xmlsearch.yandex.com/xmlsearch"):
        try:
            resp = session.get(endpoint, params=params, timeout=15)
            body = resp.text or ""
            low = body.lower()

            if "<error" in low:
                # API-level error (quota/auth/etc.)
                if "code=\"4001\"" in low or "old authorization type" in low:
                    raise RuntimeError(
                        "Yandex XML auth type error (4001): use YANDEX_XML_APIKEY for new authorization type"
                    )
                if "invalid key" in low or "forbidden" in low or "not registered" in low:
                    raise RuntimeError(
                        "Yandex XML auth error: check YANDEX_XML_APIKEY or YANDEX_XML_USER / YANDEX_XML_KEY"
                    )
                raise RuntimeError("Yandex XML returned API error; check credentials/quota in XML Search settings")

            if resp.status_code >= 400:
                resp.raise_for_status()
            time.sleep(random.uniform(*delay_range))
            return body
        except Exception as exc:
            last_error = exc
            continue

    if last_error:
        raise RuntimeError(f"Yandex XML request failed: {last_error}")

    time.sleep(random.uniform(*delay_range))
    return ""


def fetch_yandex_html_with_session(session: requests.Session, query: str, delay_range=(1.0, 2.0)) -> str:
    """Fetch Yandex SERP HTML via reusable session."""
    params = {
        "text": query,
        "lr": "213",
        "p": "0",
        "noreask": "1",
    }
    resp = session.get("https://yandex.ru/search/", params=params, timeout=15)
    resp.raise_for_status()
    time.sleep(random.uniform(*delay_range))
    return resp.text


def parse_organic_urls(html: str, limit: int = 6) -> List[str]:
    """Parse first organic URLs from Yandex SERP HTML, skipping Yandex hosts."""
    soup = BeautifulSoup(html, "html.parser")
    if is_captcha_page(html):
        return []

    seen = set()
    results: List[str] = []

    # Try multiple selector strategies because Yandex markup changes often.
    selector_groups = [
        ["li.serp-item a.Link", "li.serp-item a.OrganicTitle-Link"],
        ["div.organic a.Link", "div.organic a.OrganicTitle-Link"],
        ["article.serp-item a.Link", "article.serp-item a.OrganicTitle-Link"],
        ["a[href^='https://yandex.ru/turbo']", "a[href*='yabs.yandex']", "a[href*='yandex.ru/clck']"],
    ]

    for selectors in selector_groups:
        for selector in selectors:
            for a in soup.select(selector):
                href = a.get("href")
                if not href:
                    continue
                href = unwrap_yandex_redirect(href)
                domain = extract_domain(href)
                if not domain or domain.endswith("yandex.ru"):
                    continue
                if domain in seen:
                    continue
                seen.add(domain)
                results.append(href)
                if len(results) >= limit:
                    return results[:limit]

    # Fallback: grab any external http(s) links inside serp items.
    if len(results) < limit:
        for a in soup.select("li.serp-item a[href^='http']"):
            href = a.get("href")
            if not href:
                continue
            href = unwrap_yandex_redirect(href)
            domain = extract_domain(href)
            if not domain or domain.endswith("yandex.ru"):
                continue
            if domain in seen:
                continue
            seen.add(domain)
            results.append(href)
            if len(results) >= limit:
                break

    return results[:limit]


def parse_xml_urls(xml_text: str, limit: int = 6) -> List[str]:
    """Parse URLs from Yandex XML Search response."""
    try:
        soup = BeautifulSoup(xml_text, "xml")
    except Exception:
        # Fallback parser when lxml is not installed.
        soup = BeautifulSoup(xml_text, "html.parser")
    results: List[str] = []
    seen = set()

    for node in soup.select("group doc url"):
        url = (node.text or "").strip()
        if not url:
            continue
        domain = extract_domain(url)
        if not domain or domain.endswith("yandex.ru"):
            continue
        if domain in seen:
            continue
        seen.add(domain)
        results.append(url)
        if len(results) >= limit:
            break
    return results[:limit]


def extract_domain(url: str) -> str:
    """Extract normalized domain (strip www)."""
    try:
        host = urlparse(url).netloc.lower()
        if host.startswith("www."):
            host = host[4:]
        return host
    except Exception:
        return ""


def is_captcha_page(text: str) -> bool:
    """Return True when response appears to be a Yandex captcha challenge."""
    low = text.lower()
    markers = [
        "showcaptcha",
        "smartcaptcha",
        "checkcaptcha",
        "tmgrdfrend",
        "javascript required",
        "i am not a robot",
    ]
    return any(m in low for m in markers)


def unwrap_yandex_redirect(url: str) -> str:
    """Unwrap common Yandex redirect URLs to the real destination."""
    try:
        parsed = urlparse(url)
        host = parsed.netloc.lower()
        if "yabs.yandex" in host or "yandex.ru/clck" in url or "yandex.ru/turbo" in host:
            # Check query params like url=, u=, text=
            from urllib.parse import parse_qs

            qs = parse_qs(parsed.query)
            for key in ("url", "u", "text"):
                if key in qs and qs[key]:
                    return qs[key][0]
        return url
    except Exception:
        return url


def classify_domain(domain: str) -> str:
    """Map domain to source_type based on predefined buckets."""
    for override_domain, override_type in DOMAIN_OVERRIDES.items():
        if domain == override_domain or domain.endswith(f".{override_domain}"):
            return override_type

    for label, domains in CLASS_MAP.items():
        for base_domain in domains:
            if domain == base_domain or domain.endswith(f".{base_domain}"):
                return label
    return "other"


def suggest_source_type(domain: str) -> str:
    """Suggest source type for currently unmapped domains (non-binding)."""
    d = domain.lower()
    marketplace_markers = ["market", "shop", "store", "mall", "apteka", "pharm", "detmir"]
    reviews_markers = ["review", "otzyv", "irecommend", "otzovik"]
    media_markers = ["news", "lenta", "rbc", "kp", "vc", "journal", "mag"]
    forum_markers = ["forum", "community", "pikabu", "reddit", "babyblog"]

    if any(m in d for m in marketplace_markers):
        return "marketplace"
    if any(m in d for m in reviews_markers):
        return "reviews"
    if any(m in d for m in media_markers):
        return "media"
    if any(m in d for m in forum_markers):
        return "forum"
    return "other"


def validate_query_pool(queries: List[Tuple[str, str]]) -> None:
    """Validate query list size and category balance before running expensive collection."""
    total = len(queries)
    if total != EXPECTED_QUERY_COUNT:
        print(f"[preflight] Query count check: {total} (expected {EXPECTED_QUERY_COUNT})")
    else:
        print(f"[preflight] Query count check: {total}/{EXPECTED_QUERY_COUNT} OK")

    counts = Counter(category for _, category in queries)
    for category, expected in EXPECTED_CATEGORY_COUNTS.items():
        actual = counts.get(category, 0)
        status = "OK" if actual == expected else "MISMATCH"
        print(f"[preflight] Category {category}: {actual}/{expected} {status}")


def run_preflight(mode: str) -> None:
    """Check credentials + make one lightweight API call to ensure collection will work."""
    print("[preflight] Starting runtime checks...")
    print(f"[preflight] SEARCH_MODE={mode}")
    print(f"[preflight] YANDEX_XML_APIKEY set={bool(YANDEX_XML_APIKEY)}")
    print(f"[preflight] YANDEX_XML_USER set={bool(YANDEX_XML_USER)}")
    print(f"[preflight] YANDEX_XML_KEY set={bool(YANDEX_XML_KEY)}")

    validate_query_pool(QUERIES)

    if mode != "xml":
        print("[preflight] Skipped API probe (only required for SEARCH_MODE=xml)")
        return

    session = create_http_session()
    try:
        xml_text = fetch_yandex_xml(
            session,
            query="лучший шампунь от перхоти",
            user=YANDEX_XML_USER,
            key=YANDEX_XML_KEY,
            apikey=YANDEX_XML_APIKEY,
            delay_range=(0.0, 0.0),
        )
        urls = parse_xml_urls(xml_text, limit=3)
        if urls:
            print(f"[preflight] API probe OK, parsed {len(urls)} urls")
        else:
            print("[preflight] API probe returned response but no urls")
    except Exception as exc:
        raise SystemExit(f"[preflight] API probe failed: {exc}")


def build_dataset(queries: Iterable[Tuple[str, str]], verbose: bool = False, mode: str = "auto") -> pd.DataFrame:
    """Collect results for queries and return dataframe."""
    query_list = list(queries)
    if MAX_QUERIES > 0:
        query_list = query_list[:MAX_QUERIES]
    rows = []
    failed_queries: List[str] = []
    captcha_hits = 0
    xml_hits = 0
    session = create_http_session()

    for query, category in query_list:
        urls: List[str] = []

        if mode in ("auto", "html"):
            html = fetch_yandex_html_with_session(session, query)
            html_captcha = is_captcha_page(html)
            if html_captcha:
                captcha_hits += 1
            else:
                urls = parse_organic_urls(html, limit=6)

        has_xml_auth = bool(YANDEX_XML_APIKEY or (YANDEX_XML_USER and YANDEX_XML_KEY))
        if not urls and mode in ("auto", "xml") and has_xml_auth:
            try:
                xml_text = fetch_yandex_xml(
                    session,
                    query,
                    YANDEX_XML_USER,
                    YANDEX_XML_KEY,
                    YANDEX_XML_APIKEY,
                )
                urls = parse_xml_urls(xml_text, limit=6)
                if urls:
                    xml_hits += 1
                if urls and verbose:
                    print(f"[info] Used XML fallback for query: {query}")
            except Exception as exc:
                if verbose:
                    print(f"[warn] XML fallback failed for query '{query}': {exc}")

        if not urls:
            failed_queries.append(query)
            if verbose:
                print(f"[warn] No results parsed for query: {query}")

        for rank, url in enumerate(urls, start=1):
            domain = extract_domain(url)
            source_type = classify_domain(domain)
            rows.append(
                {
                    "query": query,
                    "category": category,
                    "url": url,
                    "domain": domain,
                    "source_type": source_type,
                    "source_rank": rank,
                }
            )

    if failed_queries:
        print(f"[summary] No results for {len(failed_queries)} of {len(query_list)} queries.")
        preview = failed_queries[:8]
        if preview:
            print("[summary] Example failed queries:")
            for q in preview:
                print(f"  - {q}")

    if captcha_hits:
        print(
            f"[summary] Captcha pages detected for {captcha_hits} queries. "
            "Direct HTML scraping is blocked on this IP."
        )
        if not (YANDEX_XML_USER and YANDEX_XML_KEY):
            print(
                "[summary] Set YANDEX_XML_APIKEY (preferred) or YANDEX_XML_USER+YANDEX_XML_KEY to enable XML fallback collection."
            )
            print("[summary] Or set YANDEX_COOKIE from a real browser session to try HTML mode.")
            print("[summary] Example PowerShell: $env:YANDEX_COOKIE='yandexuid=...; Session_id=...'")
    if xml_hits:
        print(f"[summary] XML fallback succeeded for {xml_hits} queries.")

    return pd.DataFrame(rows)


def validate_runtime_config(mode: str) -> None:
    """Validate startup config and fail fast for invalid setup."""
    allowed = {"auto", "html", "xml"}
    if mode not in allowed:
        raise SystemExit(f"Invalid SEARCH_MODE='{mode}'. Use one of: auto, html, xml")

    has_xml_auth = bool(YANDEX_XML_APIKEY or (YANDEX_XML_USER and YANDEX_XML_KEY))
    if mode == "xml" and not has_xml_auth:
        raise SystemExit(
            "SEARCH_MODE=xml requires YANDEX_XML_APIKEY or YANDEX_XML_USER and YANDEX_XML_KEY in environment."
        )


def ai_visibility_index(df: pd.DataFrame) -> pd.Series:
    """AI Visibility Index: domain_mentions / total_sources."""
    total = len(df)
    if total == 0:
        return pd.Series(dtype=float)
    counts = df["domain"].value_counts()
    return counts / total


def detect_brands_from_text(text: str) -> List[str]:
    """Detect brands from plain text using lightweight regex aliases."""
    haystack = (text or "").lower()
    found: List[str] = []
    for brand, patterns in BRAND_ALIASES.items():
        if any(re.search(pattern, haystack) for pattern in patterns):
            found.append(brand)
    return sorted(set(found))


def build_team_sheet(df: pd.DataFrame) -> pd.DataFrame:
    """Build an enriched table aligned with hackathon collaboration needs."""
    if df.empty:
        return pd.DataFrame(
            columns=[
                "query",
                "category",
                "source",
                "domain",
                "type",
                "source_rank",
                "brand_mention",
                "pg_brand",
                "pg_brand_hit",
            ]
        )

    enriched = df.copy()
    enriched["source"] = enriched.get("url", enriched["domain"]) if "url" in enriched.columns else enriched["domain"]
    enriched["type"] = enriched["source_type"]

    def _extract_mentions(row: pd.Series) -> str:
        text = (
            f"{row.get('query', '')} "
            f"{row.get('domain', '')} "
            f"{row.get('url', '')} "
            f"{row.get('ai_answer_text', '')}"
        )
        mentions = detect_brands_from_text(text)
        return ", ".join(mentions) if mentions else "none"

    enriched["brand_mention"] = enriched.apply(_extract_mentions, axis=1)
    enriched["pg_brand"] = enriched["category"].map(PG_BRAND_FOCUS).fillna("none")
    enriched["pg_brand_hit"] = enriched["brand_mention"].apply(
        lambda x: "yes" if any(pg in x for pg in PG_BRANDS) else "no"
    )
    return enriched[
        [
            "query",
            "category",
            "source",
            "domain",
            "type",
            "source_rank",
            "brand_mention",
            "pg_brand",
            "pg_brand_hit",
        ]
    ]


def save_supporting_outputs(df: pd.DataFrame) -> None:
    """Save additional summary tables useful for presentation and QA."""
    if df.empty:
        pd.DataFrame(columns=["query", "category", "sources_collected"]).to_csv(
            "query_coverage.csv", index=False, encoding="utf-8"
        )
        pd.DataFrame(columns=["domain", "mentions", "ai_visibility_index"]).to_csv(
            "domain_summary.csv", index=False, encoding="utf-8"
        )
        pd.DataFrame(columns=["domain", "mentions", "suggested_type"]).to_csv(
            "other_domains_review.csv", index=False, encoding="utf-8"
        )
        return

    coverage = (
        df.groupby(["query", "category"], as_index=False)
        .agg(sources_collected=("url", "count"))
        .sort_values(["category", "query"])
    )
    coverage.to_csv("query_coverage.csv", index=False, encoding="utf-8")

    domain_counts = df["domain"].value_counts().rename_axis("domain").reset_index(name="mentions")
    domain_counts["ai_visibility_index"] = domain_counts["mentions"] / len(df)
    domain_counts.to_csv("domain_summary.csv", index=False, encoding="utf-8")

    other_df = df[df["source_type"] == "other"]
    if other_df.empty:
        pd.DataFrame(columns=["domain", "mentions", "suggested_type"]).to_csv(
            "other_domains_review.csv", index=False, encoding="utf-8"
        )
    else:
        other_counts = other_df["domain"].value_counts().rename_axis("domain").reset_index(name="mentions")
        other_counts["suggested_type"] = other_counts["domain"].apply(suggest_source_type)
        other_counts.to_csv("other_domains_review.csv", index=False, encoding="utf-8")


def print_quality_report(df: pd.DataFrame) -> None:
    """Print compact quality diagnostics for completeness and usefulness of data."""
    print("\n=== Data Quality Check ===")
    if df.empty:
        print("rows=0, query_coverage=0.0%, avg_sources_per_query=0.00")
        return

    unique_queries = df["query"].nunique()
    total_queries = len(QUERIES) if MAX_QUERIES <= 0 else min(len(QUERIES), MAX_QUERIES)
    coverage = unique_queries / max(total_queries, 1)
    avg_sources = len(df) / max(unique_queries, 1)
    duplicates = int(df.duplicated(subset=["query", "url"]).sum())
    print(f"rows={len(df)}")
    print(f"queries_with_sources={unique_queries}/{total_queries} ({coverage:.1%})")
    print(f"avg_sources_per_query={avg_sources:.2f}")
    print(f"duplicate_query_url_rows={duplicates}")


def analytics(df: pd.DataFrame) -> None:
    """Print analytics and plot source type distribution."""
    if df.empty:
        print("[info] No data collected. Nothing to analyze.")
        return

    print("=== Source type counts ===")
    print(df["source_type"].value_counts())

    print("\n=== Unique domains ===")
    print(df["domain"].nunique())

    print("\n=== Top domains ===")
    print(df["domain"].value_counts().head(10))

    visibility = ai_visibility_index(df)
    if not visibility.empty:
        print("\n=== AI Visibility Index (top 10) ===")
        print(visibility.head(10))

    ax = df["source_type"].value_counts().plot(kind="bar", title="Source Type Distribution")
    ax.set_xlabel("Source Type")
    ax.set_ylabel("Count")
    plt.tight_layout()
    plt.savefig("source_type_distribution.png", dpi=160)
    if PLOT_SHOW:
        plt.show()
    else:
        plt.close()


def main():
    validate_runtime_config(SEARCH_MODE)
    run_preflight(SEARCH_MODE)
    if VERBOSE_LOGS:
        print(f"[info] SEARCH_MODE={SEARCH_MODE}")
    df = build_dataset(QUERIES, verbose=VERBOSE_LOGS, mode=SEARCH_MODE)
    # Required dataset for core analysis.
    if df.empty:
        pd.DataFrame(columns=OUTPUT_COLUMNS).to_csv("sources.csv", index=False, encoding="utf-8")
    else:
        df[OUTPUT_COLUMNS].to_csv("sources.csv", index=False, encoding="utf-8")

    # Additional outputs for team workflow and presentation prep.
    team_df = build_team_sheet(df)
    team_df.to_csv("sources_team_sheet.csv", index=False, encoding="utf-8")
    save_supporting_outputs(df)

    print(f"Saved {len(df)} rows to sources.csv")
    print("Saved sources_team_sheet.csv, query_coverage.csv, domain_summary.csv")
    print_quality_report(df)
    analytics(df)


if __name__ == "__main__":
    main()