"""
Daily aggregation job: reads raw 5-min data from vm_metrics,
computes hourly aggregates, and writes to vm_metrics_hourly.

Idempotent: only processes dates not yet present in the hourly tables.
Run daily via cron (run_aggregate.sh).
"""

import logging
import os
import sys

import toml
import yaml
from google.cloud import bigquery
from google.oauth2 import service_account

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
)
logger = logging.getLogger("vm_aggregate")

PROJECT_ID = "gcp-digital-bi-dev"
RAW_DATASET = "vm_metrics"
HOURLY_DATASET = "vm_metrics_hourly"

FIXED_COLUMNS = frozenset({
    "metric_timestamp", "metric_value", "metric_labels_hash",
    "metric_date", "kubernetes", "provider", "cloud_region",
    "_dlt_load_id", "_dlt_id",
})


def get_bq_client() -> bigquery.Client:
    secrets_path = os.path.join(os.path.dirname(__file__), ".dlt", "secrets.toml")
    with open(secrets_path) as f:
        secrets = toml.load(f)

    creds_info = dict(secrets["destination"]["bigquery"]["credentials"])
    creds_info["type"] = "service_account"
    creds_info["token_uri"] = "https://oauth2.googleapis.com/token"

    credentials = service_account.Credentials.from_service_account_info(creds_info)
    return bigquery.Client(project=PROJECT_ID, credentials=credentials)


def ensure_hourly_dataset(client: bigquery.Client) -> None:
    dataset_ref = bigquery.DatasetReference(PROJECT_ID, HOURLY_DATASET)
    dataset = bigquery.Dataset(dataset_ref)
    dataset.location = "us-west2"
    client.create_dataset(dataset, exists_ok=True)


def get_raw_tables(client: bigquery.Client) -> list[str]:
    """Return metric table names from the raw dataset (exclude dlt internal tables)."""
    query = f"""
        SELECT table_name
        FROM `{PROJECT_ID}.{RAW_DATASET}.INFORMATION_SCHEMA.TABLES`
        WHERE table_name NOT LIKE '_dlt%'
        ORDER BY table_name
    """
    return [row.table_name for row in client.query(query).result()]


def get_label_columns(client: bigquery.Client, table_name: str) -> list[str]:
    """Discover dynamic label columns (everything that's not a fixed column)."""
    query = f"""
        SELECT column_name
        FROM `{PROJECT_ID}.{RAW_DATASET}.INFORMATION_SCHEMA.COLUMNS`
        WHERE table_name = '{table_name}'
        ORDER BY ordinal_position
    """
    return [
        row.column_name
        for row in client.query(query).result()
        if row.column_name not in FIXED_COLUMNS
    ]


def get_unaggregated_dates(
    client: bigquery.Client, table_name: str,
) -> list[str]:
    """Find dates in raw that don't yet exist in hourly."""
    hourly_table = f"{PROJECT_ID}.{HOURLY_DATASET}.{table_name}"

    hourly_exists = True
    try:
        client.get_table(hourly_table)
    except Exception:
        hourly_exists = False

    if hourly_exists:
        query = f"""
            SELECT DISTINCT metric_date
            FROM `{PROJECT_ID}.{RAW_DATASET}.{table_name}`
            WHERE metric_date NOT IN (
                SELECT DISTINCT DATE(hour_timestamp) FROM `{hourly_table}`
            )
            ORDER BY metric_date
        """
    else:
        query = f"""
            SELECT DISTINCT metric_date
            FROM `{PROJECT_ID}.{RAW_DATASET}.{table_name}`
            ORDER BY metric_date
        """

    return [str(row.metric_date) for row in client.query(query).result()]


def build_aggregation_sql(
    table_name: str,
    label_columns: list[str],
    target_date: str,
) -> str:
    label_selects = ",\n        ".join(
        f"ANY_VALUE({col}) AS {col}" for col in label_columns
    )
    if label_selects:
        label_selects = ",\n        " + label_selects

    return f"""
    SELECT
        metric_labels_hash,
        kubernetes,
        provider,
        cloud_region,
        TIMESTAMP_SECONDS(
            metric_timestamp - MOD(metric_timestamp, 3600)
        ) AS hour_timestamp,
        MIN(metric_value) AS value_min,
        MAX(metric_value) AS value_max,
        AVG(metric_value) AS value_avg,
        ARRAY_AGG(metric_value ORDER BY metric_timestamp ASC LIMIT 1)[OFFSET(0)] AS value_first,
        ARRAY_AGG(metric_value ORDER BY metric_timestamp DESC LIMIT 1)[OFFSET(0)] AS value_last,
        COUNT(*) AS sample_count{label_selects}
    FROM `{PROJECT_ID}.{RAW_DATASET}.{table_name}`
    WHERE metric_date = '{target_date}'
    GROUP BY
        metric_labels_hash,
        kubernetes,
        provider,
        cloud_region,
        hour_timestamp
    """


def create_hourly_table_if_needed(
    client: bigquery.Client,
    table_name: str,
    label_columns: list[str],
) -> None:
    """Create the hourly table with partitioning and clustering if it doesn't exist."""
    table_ref = f"{PROJECT_ID}.{HOURLY_DATASET}.{table_name}"
    try:
        client.get_table(table_ref)
        return
    except Exception:
        pass

    schema = [
        bigquery.SchemaField("metric_labels_hash", "STRING"),
        bigquery.SchemaField("kubernetes", "STRING"),
        bigquery.SchemaField("provider", "STRING"),
        bigquery.SchemaField("cloud_region", "STRING"),
        bigquery.SchemaField("hour_timestamp", "TIMESTAMP"),
        bigquery.SchemaField("value_min", "FLOAT64"),
        bigquery.SchemaField("value_max", "FLOAT64"),
        bigquery.SchemaField("value_avg", "FLOAT64"),
        bigquery.SchemaField("value_first", "FLOAT64"),
        bigquery.SchemaField("value_last", "FLOAT64"),
        bigquery.SchemaField("sample_count", "INT64"),
    ]
    for col in label_columns:
        schema.append(bigquery.SchemaField(col, "STRING"))

    table = bigquery.Table(table_ref, schema=schema)
    table.time_partitioning = bigquery.TimePartitioning(
        type_=bigquery.TimePartitioningType.DAY,
        field="hour_timestamp",
    )
    table.clustering_fields = ["kubernetes", "metric_labels_hash"]

    client.create_table(table)
    logger.info("Created hourly table: %s", table_ref)


def aggregate_table(client: bigquery.Client, table_name: str) -> None:
    label_columns = get_label_columns(client, table_name)
    dates = get_unaggregated_dates(client, table_name)

    if not dates:
        logger.info("[%s] No unaggregated dates, skipping", table_name)
        return

    logger.info("[%s] %d dates to aggregate: %s .. %s", table_name, len(dates), dates[0], dates[-1])

    create_hourly_table_if_needed(client, table_name, label_columns)

    hourly_table = f"{PROJECT_ID}.{HOURLY_DATASET}.{table_name}"

    for target_date in dates:
        agg_sql = build_aggregation_sql(table_name, label_columns, target_date)
        insert_sql = f"INSERT INTO `{hourly_table}`\n{agg_sql}"

        logger.info("[%s] Aggregating %s ...", table_name, target_date)
        job = client.query(insert_sql)
        job.result()
        logger.info(
            "[%s] %s done: %s bytes processed",
            table_name, target_date, job.total_bytes_processed,
        )


def run_aggregation() -> None:
    client = get_bq_client()
    ensure_hourly_dataset(client)

    tables = get_raw_tables(client)
    if not tables:
        logger.warning("No raw tables found in %s.%s", PROJECT_ID, RAW_DATASET)
        return

    logger.info("Found %d raw tables to aggregate", len(tables))

    for table_name in tables:
        try:
            aggregate_table(client, table_name)
        except Exception:
            logger.exception("Failed to aggregate %s, continuing", table_name)


if __name__ == "__main__":
    run_aggregation()