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")
    for row in alias_rows:
        try:
            from google.cloud import bigquery

            query = f"""
                MERGE `{alias_table}` T
                USING (
                    SELECT
                        @fg_id AS fg_id,
                        @id_system AS id_system,
                        @id_value AS id_value,
                        @alias_key AS alias_key,
                        @is_primary AS is_primary,
                        @confidence AS confidence,
                        @source AS source,
                        @job_id AS job_id,
                        @entity_type AS entity_type
                ) 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.ScalarQueryParameter("fg_id", "STRING", row.get("fg_id")),
                bigquery.ScalarQueryParameter("id_system", "STRING", row.get("id_system")),
                bigquery.ScalarQueryParameter("id_value", "STRING", row.get("id_value")),
                bigquery.ScalarQueryParameter("alias_key", "STRING", row.get("alias_key")),
                bigquery.ScalarQueryParameter("is_primary", "BOOL", row.get("is_primary")),
                bigquery.ScalarQueryParameter("confidence", "FLOAT64", row.get("confidence")),
                bigquery.ScalarQueryParameter("source", "STRING", row.get("source")),
                bigquery.ScalarQueryParameter("job_id", "STRING", row.get("job_id")),
                bigquery.ScalarQueryParameter("entity_type", "STRING", row.get("entity_type")),
            ]
            _run_query(query, params)
        except Exception as exc:
            logger.warning("Alias merge failed for fg_id=%s: %s", row.get("fg_id"), 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")
    for row in event_rows:
        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)
                VALUES (@event_id, @subject_fg_id, @event_type, @merged_from_fg_id, @merged_into_fg_id, @reason, @actor, @job_id)
            """
            params = [
                bigquery.ScalarQueryParameter("event_id", "STRING", row.get("event_id")),
                bigquery.ScalarQueryParameter("subject_fg_id", "STRING", row.get("subject_fg_id")),
                bigquery.ScalarQueryParameter("event_type", "STRING", row.get("event_type")),
                bigquery.ScalarQueryParameter("merged_from_fg_id", "STRING", row.get("merged_from_fg_id")),
                bigquery.ScalarQueryParameter("merged_into_fg_id", "STRING", row.get("merged_into_fg_id")),
                bigquery.ScalarQueryParameter("reason", "STRING", row.get("reason")),
                bigquery.ScalarQueryParameter("actor", "STRING", row.get("actor")),
                bigquery.ScalarQueryParameter("job_id", "STRING", row.get("job_id")),
            ]
            _run_query(query, params)
        except Exception as exc:
            logger.warning("Event insert failed for fg_id=%s: %s", row.get("subject_fg_id"), 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,
                    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 -> id_system='domain'
    - From company_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(dataset, "company_external_ids")

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

    try:
        # Pull core identifiers from companies_gold
        from google.cloud import bigquery

        query = f"""
            SELECT fg_id, wikidata_id, domain
            FROM `{companies_table}`
            WHERE fg_id IS NOT NULL
        """
        rows = _run_query(query)
        alias_rows: List[Dict[str, Any]] = []
        for row in rows:
            fg_id = getattr(row, "fg_id", None)
            wikidata_id = getattr(row, "wikidata_id", None)
            domain = getattr(row, "domain", None)
            if wikidata_id:
                norm = AliasNormalizer.normalize("wikidata", wikidata_id)
                alias_rows.append(
                    {
                        "fg_id": fg_id,
                        "id_system": "wikidata",
                        "id_value": norm,
                        "alias_key": AliasNormalizer.compute_alias_key("wikidata", norm),
                        "is_primary": True,
                        "confidence": None,
                        "source": "backfill",
                        "job_id": job_id,
                        "entity_type": "company",
                    }
                )
            if domain:
                norm = AliasNormalizer.normalize("domain", domain)
                if norm:
                    alias_rows.append(
                        {
                            "fg_id": fg_id,
                            "id_system": "domain",
                            "id_value": norm,
                            "alias_key": AliasNormalizer.compute_alias_key("domain", norm),
                            "is_primary": False,
                            "confidence": None,
                            "source": "backfill",
                            "job_id": job_id,
                            "entity_type": "company",
                        }
                    )

        _merge_alias_rows(alias_rows, dataset)
        processed += len(alias_rows)
    except Exception as exc:
        errors.append(str(exc))
        logger.warning("populate_aliases_from_columns base load failed: %s", exc)

    try:
        # External IDs (ticker/lei/isin) joined to fg_id
        from google.cloud import bigquery

        external_query = f"""
            SELECT cg.fg_id AS fg_id, ce.id_type AS id_type, ce.id_value AS id_value
            FROM `{external_ids_table}` ce
            JOIN `{companies_table}` cg
              ON ce.wikidata_id = cg.wikidata_id
            WHERE cg.fg_id IS NOT NULL
              AND ce.id_type IN ('ticker', 'lei', 'isin', 'duns', 'sec_cik')
        """
        rows = _run_query(external_query)
        alias_rows = []
        for row in rows:
            fg_id = getattr(row, "fg_id", None)
            id_type = getattr(row, "id_type", None)
            id_value = getattr(row, "id_value", None)
            if not (fg_id and id_type and id_value):
                continue
            norm = AliasNormalizer.normalize(id_type, id_value)
            if not norm:
                continue
            alias_rows.append(
                {
                    "fg_id": fg_id,
                    "id_system": id_type,
                    "id_value": norm,
                    "alias_key": AliasNormalizer.compute_alias_key(id_type, norm),
                    "is_primary": False,
                    "confidence": None,
                    "source": "backfill",
                    "job_id": job_id,
                    "entity_type": "company",
                }
            )
            if len(alias_rows) >= batch_size:
                _merge_alias_rows(alias_rows, dataset)
                processed += len(alias_rows)
                alias_rows = []

        if alias_rows:
            _merge_alias_rows(alias_rows, dataset)
            processed += len(alias_rows)
    except Exception as exc:
        errors.append(str(exc))
        logger.warning("populate_aliases_from_columns external ids failed: %s", exc)

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


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,
    }
