from __future__ import annotations

import logging
import os
from typing import Any, Dict, List, Optional

from .fg_entity_id import generate_event_id, generate_fg_id
from .fg_entity_registry import AliasNormalizer

logger = logging.getLogger(__name__)

_bq_client = None


def _project_id() -> str:
    return (
        os.getenv("FM_FOUNDRYGRAPH_PROJECT_ID")
        or os.getenv("GOOGLE_CLOUD_PROJECT")
        or ""
    ).strip()


def _table_ref(dataset: str, table: str) -> str:
    project_id = _project_id()
    if project_id:
        return f"{project_id}.{dataset}.{table}"
    return f"{dataset}.{table}"


def _get_bq_client():
    global _bq_client
    project_id = _project_id()
    if _bq_client is None or (project_id and not getattr(_bq_client, "project", None)):
        from google.cloud import bigquery

        _bq_client = bigquery.Client(project=project_id) if project_id else bigquery.Client()
    return _bq_client


def _run_query(query: str, params: Optional[List[Any]] = None) -> List[Any]:
    client = _get_bq_client()
    from google.cloud import bigquery

    job_config = bigquery.QueryJobConfig(query_parameters=params or [])
    return list(client.query(query, job_config=job_config, project=_project_id() or None).result())


def _merge_alias_rows(alias_rows: List[Dict[str, Any]], dataset: str) -> None:
    if not alias_rows:
        return
    alias_table = _table_ref(dataset, "company_aliases")
    try:
        from google.cloud import bigquery

        query = f"""
            MERGE `{alias_table}` T
            USING UNNEST(@alias_rows) AS S
            ON T.alias_key = S.alias_key
            WHEN NOT MATCHED THEN
              INSERT (fg_id, id_system, id_value, alias_key, is_primary, confidence, source, job_id, entity_type)
              VALUES (S.fg_id, S.id_system, S.id_value, S.alias_key, S.is_primary, S.confidence, S.source, S.job_id, S.entity_type)
        """
        params = [
            bigquery.ArrayQueryParameter(
                "alias_rows",
                "STRUCT<fg_id STRING, id_system STRING, id_value STRING, alias_key STRING, is_primary BOOL, confidence FLOAT64, source STRING, job_id STRING, entity_type STRING>",
                alias_rows,
            )
        ]
        _run_query(query, params)
    except Exception as exc:
        logger.warning("Alias merge failed for %s rows: %s", len(alias_rows), exc)


def _insert_event_rows(event_rows: List[Dict[str, Any]], dataset: str) -> None:
    if not event_rows:
        return
    events_table = _table_ref(dataset, "company_events")
    try:
        from google.cloud import bigquery

        query = f"""
            INSERT `{events_table}`
            (event_id, subject_fg_id, event_type, merged_from_fg_id, merged_into_fg_id, reason, actor, job_id)
            SELECT
              event_id,
              subject_fg_id,
              event_type,
              merged_from_fg_id,
              merged_into_fg_id,
              reason,
              actor,
              job_id
            FROM UNNEST(@event_rows)
        """
        params = [
            bigquery.ArrayQueryParameter(
                "event_rows",
                "STRUCT<event_id STRING, subject_fg_id STRING, event_type STRING, merged_from_fg_id STRING, merged_into_fg_id STRING, reason STRING, actor STRING, job_id STRING>",
                event_rows,
            )
        ]
        _run_query(query, params)
    except Exception as exc:
        logger.warning("Event insert failed for %s rows: %s", len(event_rows), exc)


async def backfill_existing_companies(
    batch_size: int = 10000,
    dry_run: bool = True,
    job_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    One-time backfill - IDEMPOTENT and RESUMABLE.

    Strategy:
    1. Query companies_gold WHERE fg_id IS NULL ORDER BY wikidata_id
    2. Process in deterministic batches (ordered by PK for resumability)
    3. For each row:
       - Generate fg_id
       - Set fg_id_created_at = CURRENT_TIMESTAMP()
       - Insert into company_aliases (wikidata as primary)
       - Log 'created' event with job_id
    4. Use MERGE/UPDATE (not replace table)
    5. Only assign fg_id where fg_id IS NULL (skip already-filled)

    Returns: {total: int, processed: int, skipped: int, errors: list}
    """
    errors: List[str] = []
    dataset = "foundrygraph_gold"
    companies_table = _table_ref(dataset, "companies_gold")

    try:
        if dry_run:
            count_query = f"""
                SELECT COUNTIF(fg_id IS NULL) AS missing_fg_id
                FROM `{companies_table}`
            """
            rows = _run_query(count_query)
            missing = int(getattr(rows[0], "missing_fg_id", 0)) if rows else 0
            return {
                "total": missing,
                "processed": 0,
                "skipped": missing,
                "errors": [],
            }
    except Exception as exc:
        return {
            "total": 0,
            "processed": 0,
            "skipped": 0,
            "errors": [str(exc)],
        }

    processed = 0
    total_missing = 0
    while True:
        try:
            from google.cloud import bigquery

            select_query = f"""
                SELECT wikidata_id
                FROM `{companies_table}`
                WHERE fg_id IS NULL AND wikidata_id IS NOT NULL
                ORDER BY wikidata_id
                LIMIT @limit
            """
            rows = _run_query(
                select_query,
                [bigquery.ScalarQueryParameter("limit", "INT64", int(batch_size))],
            )
        except Exception as exc:
            errors.append(str(exc))
            break

        if not rows:
            break

        total_missing += len(rows)
        updates: List[Dict[str, Any]] = []
        alias_rows: List[Dict[str, Any]] = []
        event_rows: List[Dict[str, Any]] = []

        for row in rows:
            wikidata_id = getattr(row, "wikidata_id", None)
            if not wikidata_id:
                continue
            fg_id = generate_fg_id("co")
            updates.append({"wikidata_id": wikidata_id, "fg_id": fg_id})

            normalized_wd = AliasNormalizer.normalize("wikidata", wikidata_id)
            alias_rows.append(
                {
                    "fg_id": fg_id,
                    "id_system": "wikidata",
                    "id_value": normalized_wd,
                    "alias_key": AliasNormalizer.compute_alias_key(
                        "wikidata", normalized_wd
                    ),
                    "is_primary": True,
                    "confidence": None,
                    "source": "backfill",
                    "job_id": job_id,
                    "entity_type": "company",
                }
            )

            event_rows.append(
                {
                    "event_id": generate_event_id(),
                    "subject_fg_id": fg_id,
                    "event_type": "created",
                    "merged_from_fg_id": None,
                    "merged_into_fg_id": None,
                    "reason": None,
                    "actor": "system",
                    "job_id": job_id,
                }
            )

        try:
            from google.cloud import bigquery

            updates_table = _table_ref(dataset, "companies_gold")
            update_query = f"""
                MERGE `{updates_table}` T
                USING UNNEST(@updates) AS S
                ON T.wikidata_id = S.wikidata_id
                WHEN MATCHED AND T.fg_id IS NULL THEN
                  UPDATE SET
                    fg_id = S.fg_id,
                    status = COALESCE(T.status, 'active'),
                    fg_id_created_at = CURRENT_TIMESTAMP(),
                    fg_id_updated_at = CURRENT_TIMESTAMP()
            """
            update_params = [
                bigquery.ArrayQueryParameter(
                    "updates",
                    "STRUCT<wikidata_id STRING, fg_id STRING>",
                    updates,
                )
            ]
            _run_query(update_query, update_params)

            _merge_alias_rows(alias_rows, dataset)
            _insert_event_rows(event_rows, dataset)
            processed += len(updates)
        except Exception as exc:
            errors.append(str(exc))
            logger.warning("Backfill batch failed: %s", exc)
            break

    skipped = max(0, total_missing - processed)
    return {
        "total": total_missing,
        "processed": processed,
        "skipped": skipped,
        "errors": errors,
    }


async def populate_aliases_from_columns(
    batch_size: int = 10000,
    job_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Populate company_aliases from existing identifier columns.

    Sources:
    - wikidata_id -> id_system='wikidata', is_primary=True
    - domain/primary_domain -> id_system='domain'
    - From silver.external_ids: ticker, lei, isin

    All values normalized at write-time via AliasNormalizer.
    """
    dataset = "foundrygraph_gold"
    companies_table = _table_ref(dataset, "companies_gold")
    external_ids_table = _table_ref("foundrygraph_silver", "external_ids")

    errors: List[str] = []
    processed = 0

    try:
        from google.cloud import bigquery

        last_wikidata_id = None
        while True:
            query = f"""
                SELECT fg_id, wikidata_id, domain, primary_domain
                FROM `{companies_table}`
                WHERE fg_id IS NOT NULL
                  AND wikidata_id IS NOT NULL
                  {"AND wikidata_id > @after_wikidata_id" if last_wikidata_id else ""}
                ORDER BY wikidata_id
                LIMIT @limit
            """
            params = [bigquery.ScalarQueryParameter("limit", "INT64", int(batch_size))]
            if last_wikidata_id:
                params.insert(
                    0,
                    bigquery.ScalarQueryParameter("after_wikidata_id", "STRING", last_wikidata_id),
                )
            rows = _run_query(query, params)
            if not rows:
                break

            alias_rows: List[Dict[str, Any]] = []
            fg_id_by_wikidata: Dict[str, str] = {}
            for row in rows:
                fg_id = getattr(row, "fg_id", None)
                wikidata_id = getattr(row, "wikidata_id", None)
                if not (fg_id and wikidata_id):
                    continue

                fg_id_by_wikidata[wikidata_id] = fg_id
                wikidata_norm = AliasNormalizer.normalize("wikidata", wikidata_id)
                alias_rows.append(
                    {
                        "fg_id": fg_id,
                        "id_system": "wikidata",
                        "id_value": wikidata_norm,
                        "alias_key": AliasNormalizer.compute_alias_key("wikidata", wikidata_norm),
                        "is_primary": True,
                        "confidence": None,
                        "source": "backfill",
                        "job_id": job_id,
                        "entity_type": "company",
                    }
                )

                candidate_domains = []
                primary_domain = getattr(row, "primary_domain", None)
                domain = getattr(row, "domain", None)
                if primary_domain:
                    candidate_domains.append(primary_domain)
                if domain and domain != primary_domain:
                    candidate_domains.append(domain)
                seen_domain_norms = set()
                for domain_value in candidate_domains:
                    domain_norm = AliasNormalizer.normalize("domain", domain_value)
                    if not domain_norm or domain_norm in seen_domain_norms:
                        continue
                    seen_domain_norms.add(domain_norm)
                    alias_rows.append(
                        {
                            "fg_id": fg_id,
                            "id_system": "domain",
                            "id_value": domain_norm,
                            "alias_key": AliasNormalizer.compute_alias_key("domain", domain_norm),
                            "is_primary": False,
                            "confidence": None,
                            "source": "backfill",
                            "job_id": job_id,
                            "entity_type": "company",
                        }
                    )

            if alias_rows:
                _merge_alias_rows(alias_rows, dataset)
                processed += len(alias_rows)

            external_rows = _run_query(
                f"""
                    SELECT wikidata_id, id_type, id_value
                    FROM `{external_ids_table}`
                    WHERE wikidata_id IN UNNEST(@wikidata_ids)
                      AND LOWER(id_type) IN ('ticker', 'lei', 'isin', 'duns', 'cik', 'sec_cik')
                """,
                [
                    bigquery.ArrayQueryParameter(
                        "wikidata_ids",
                        "STRING",
                        sorted(fg_id_by_wikidata),
                    )
                ],
            )
            external_alias_rows: List[Dict[str, Any]] = []
            seen_external_alias_keys = set()
            for row in external_rows:
                wikidata_id = getattr(row, "wikidata_id", None)
                id_type = (getattr(row, "id_type", None) or "").lower()
                id_value = getattr(row, "id_value", None)
                fg_id = fg_id_by_wikidata.get(wikidata_id)
                if not (fg_id and id_type and id_value):
                    continue
                norm = AliasNormalizer.normalize(id_type, id_value)
                if not norm:
                    continue
                alias_key = AliasNormalizer.compute_alias_key(id_type, norm)
                if alias_key in seen_external_alias_keys:
                    continue
                seen_external_alias_keys.add(alias_key)
                external_alias_rows.append(
                    {
                        "fg_id": fg_id,
                        "id_system": id_type,
                        "id_value": norm,
                        "alias_key": alias_key,
                        "is_primary": False,
                        "confidence": None,
                        "source": "backfill",
                        "job_id": job_id,
                        "entity_type": "company",
                    }
                )

            if external_alias_rows:
                _merge_alias_rows(external_alias_rows, dataset)
                processed += len(external_alias_rows)

            last_wikidata_id = getattr(rows[-1], "wikidata_id", None)
    except Exception as exc:
        errors.append(str(exc))
        logger.warning("populate_aliases_from_columns failed: %s", exc)

    return {
        "processed": processed,
        "errors": errors,
    }


async def populate_events_from_fg_state(
    batch_size: int = 10000,
    job_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Backfill a baseline `created` event for entities that already have fg_ids.

    Rebuilds can restore fg_ids from a seed table without recreating the event
    log, so this step ensures `company_events` is populated for those entities.
    """
    dataset = "foundrygraph_gold"
    companies_table = _table_ref(dataset, "companies_gold")
    events_table = _table_ref(dataset, "company_events")
    errors: List[str] = []
    processed = 0

    try:
        from google.cloud import bigquery

        last_fg_id = None
        while True:
            query = f"""
                SELECT cg.fg_id
                FROM `{companies_table}` cg
                LEFT JOIN `{events_table}` ce
                  ON cg.fg_id = ce.subject_fg_id
                 AND ce.event_type = 'created'
                WHERE cg.fg_id IS NOT NULL
                  AND ce.event_id IS NULL
                  {"AND cg.fg_id > @after_fg_id" if last_fg_id else ""}
                ORDER BY cg.fg_id
                LIMIT @limit
            """
            params = [bigquery.ScalarQueryParameter("limit", "INT64", int(batch_size))]
            if last_fg_id:
                params.insert(
                    0,
                    bigquery.ScalarQueryParameter("after_fg_id", "STRING", last_fg_id),
                )
            rows = _run_query(query, params)
            if not rows:
                break

            event_rows = [
                {
                    "event_id": generate_event_id(),
                    "subject_fg_id": getattr(row, "fg_id", None),
                    "event_type": "created",
                    "merged_from_fg_id": None,
                    "merged_into_fg_id": None,
                    "reason": "baseline_backfill",
                    "actor": "system",
                    "job_id": job_id,
                }
                for row in rows
                if getattr(row, "fg_id", None)
            ]
            if event_rows:
                _insert_event_rows(event_rows, dataset)
                processed += len(event_rows)

            last_fg_id = getattr(rows[-1], "fg_id", None)
    except Exception as exc:
        errors.append(str(exc))
        logger.warning("populate_events_from_fg_state failed: %s", exc)

    return {
        "processed": processed,
        "errors": errors,
    }


async def ensure_fg_id_registry(
    batch_size: int = 10000,
    job_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Ensure canonical Gold has fg_ids, aliases, and baseline events after a rebuild/cutover.

    Runs the missing-id backfill first, then refreshes alias coverage from the
    identifier columns and external IDs that point at the now-populated fg_ids,
    and finally backfills baseline events for restored entities.
    """
    backfill = await backfill_existing_companies(
        batch_size=batch_size,
        dry_run=False,
        job_id=job_id,
    )
    aliases = await populate_aliases_from_columns(
        batch_size=batch_size,
        job_id=job_id,
    )
    events = await populate_events_from_fg_state(
        batch_size=batch_size,
        job_id=job_id,
    )
    return {
        "backfill": backfill,
        "aliases": aliases,
        "events": events,
    }


async def verify_backfill_integrity() -> Dict[str, Any]:
    """
    Verify backfill completed correctly:
    - All companies have fg_id (no NULLs)
    - All fg_ids are unique among active entities
    - All fg_ids with status='deprecated' have valid redirect_to_fg_id
    - Every fg_id has at least: domain OR wikidata_id OR one alias row

    Returns: {passed: bool, checks: dict[str, {passed, count, errors}]}
    """
    dataset = "foundrygraph_gold"
    companies_table = _table_ref(dataset, "companies_gold")
    alias_table = _table_ref(dataset, "company_aliases")

    checks: Dict[str, Dict[str, Any]] = {}
    passed = True

    # 1) NULL fg_id
    null_query = f"""
        SELECT COUNT(*) AS cnt
        FROM `{companies_table}`
        WHERE fg_id IS NULL
    """
    null_rows = _run_query(null_query)
    null_count = int(getattr(null_rows[0], "cnt", 0)) if null_rows else 0
    checks["missing_fg_id"] = {
        "passed": null_count == 0,
        "count": null_count,
        "errors": [] if null_count == 0 else ["Found companies with NULL fg_id"],
    }
    passed = passed and null_count == 0

    # 2) Uniqueness among active entities
    uniq_query = f"""
        SELECT COUNT(*) AS total, COUNT(DISTINCT fg_id) AS unique_count
        FROM `{companies_table}`
        WHERE COALESCE(status, 'active') = 'active'
    """
    uniq_rows = _run_query(uniq_query)
    total = int(getattr(uniq_rows[0], "total", 0)) if uniq_rows else 0
    unique_count = int(getattr(uniq_rows[0], "unique_count", 0)) if uniq_rows else 0
    uniq_ok = total == unique_count
    checks["unique_active_fg_id"] = {
        "passed": uniq_ok,
        "count": total - unique_count if total >= unique_count else 0,
        "errors": [] if uniq_ok else ["Duplicate fg_id among active entities"],
    }
    passed = passed and uniq_ok

    # 3) Deprecated redirect validity
    redirect_query = f"""
        SELECT COUNT(*) AS cnt
        FROM `{companies_table}` v
        LEFT JOIN `{companies_table}` s
          ON v.redirect_to_fg_id = s.fg_id
         AND COALESCE(s.status, 'active') = 'active'
        WHERE v.status = 'deprecated'
          AND (v.redirect_to_fg_id IS NULL OR s.fg_id IS NULL)
    """
    redirect_rows = _run_query(redirect_query)
    redirect_count = int(getattr(redirect_rows[0], "cnt", 0)) if redirect_rows else 0
    redirect_ok = redirect_count == 0
    checks["deprecated_redirects"] = {
        "passed": redirect_ok,
        "count": redirect_count,
        "errors": [] if redirect_ok else ["Deprecated entities missing valid redirect"],
    }
    passed = passed and redirect_ok

    # 4) Alias coverage
    coverage_query = f"""
        SELECT COUNT(*) AS cnt
        FROM `{companies_table}` cg
        LEFT JOIN `{alias_table}` ca
          ON cg.fg_id = ca.fg_id
        WHERE cg.fg_id IS NOT NULL
          AND cg.domain IS NULL
          AND cg.wikidata_id IS NULL
          AND ca.fg_id IS NULL
    """
    coverage_rows = _run_query(coverage_query)
    coverage_count = int(getattr(coverage_rows[0], "cnt", 0)) if coverage_rows else 0
    coverage_ok = coverage_count == 0
    checks["alias_coverage"] = {
        "passed": coverage_ok,
        "count": coverage_count,
        "errors": [] if coverage_ok else ["Entities missing domain/wikidata/alias coverage"],
    }
    passed = passed and coverage_ok

    return {
        "passed": passed,
        "checks": checks,
    }
