"""
dlt pipeline: TiDB Cloud Observability -> BigQuery

Incrementally queries time series from TiDB Cloud's VictoriaMetrics-backed
observability endpoints via /api/v1/query_range and loads them into BigQuery.

Each metric becomes a single BigQuery table (all endpoints append into it).
Tables are date-partitioned on metric_date and clustered by kubernetes.
"""

import logging
import os
import resource as resource_mod
import time
from typing import Any, Iterator

import yaml
import dlt
from dlt.destinations.impl.bigquery.bigquery_adapter import bigquery_adapter

from vm_common import (
    DATASET_NAME,
    fetch_query_range,
    parse_query_range_result,
)

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

dlt.config["destination.bigquery.location"] = "us-west2"


def make_resource(
    query: str,
    table_name: str,
    base_url: str,
    resource_name: str,
    kubernetes: str,
    provider: str,
    cloud_region: str,
    default_start: int,
    step: int = 300,
    window_seconds: int = 3600,
    keep_labels: list[str] | None = None,
):
    """Factory: creates a dlt resource for one metric+endpoint combination."""

    @dlt.resource(
        name=resource_name,
        table_name=table_name,
        write_disposition="append",
    )
    def _resource() -> Iterator[dict]:
        state = dlt.current.resource_state()
        start = state.setdefault("last_end", default_start)
        end = start + window_seconds

        current_time = int(time.time())
        latest_end = current_time - (current_time % window_seconds)

        if end > latest_end:
            logger.info("[%s] Caught up (end %d > latest %d), skipping", resource_name, end, latest_end)
            return

        logger.info("[%s] query_range %d -> %d (step=%d)", resource_name, start, end, step)

        try:
            results = fetch_query_range(base_url, query, start, end - 1, step)
        except Exception:
            logger.exception("[%s] Failed to fetch, skipping this window", resource_name)
            return

        row_count = 0
        for row in parse_query_range_result(results, kubernetes, provider, cloud_region, keep_labels):
            row_count += 1
            yield row

        logger.info("[%s] Yielded %d rows", resource_name, row_count)
        state["last_end"] = end

    return _resource


@dlt.source(name="victoriametrics")
def vm_source(
    metric_cfg: dict,
    endpoints: list[dict],
    default_start: int,
    window_seconds: int,
) -> Any:
    """Yields one resource per endpoint for a single metric."""
    for endpoint in endpoints:
        resource_name = f"{metric_cfg['name']}__{endpoint['kubernetes']}"
        resource = make_resource(
            query=metric_cfg["query"],
            table_name=metric_cfg["name"],
            base_url=endpoint["url"],
            resource_name=resource_name,
            kubernetes=endpoint["kubernetes"],
            provider=endpoint["provider"],
            cloud_region=endpoint["cloud_region"],
            default_start=default_start,
            step=metric_cfg.get("step", 300),
            window_seconds=window_seconds,
            keep_labels=metric_cfg.get("labels"),
        )
        bigquery_adapter(
            resource,
            partition="metric_date",
            cluster=["kubernetes", "metric_labels_hash"],
        )
        yield resource


def load_sources_config() -> dict:
    config_path = os.path.join(os.path.dirname(__file__), "sources.yml")
    with open(config_path, "r") as f:
        return yaml.safe_load(f)


MAX_RUN_SECONDS = 25 * 60  # stay within 30-min cron interval


def run_pipeline(destination: str = "bigquery") -> None:
    """Main entry point: loops through windows until caught up or time budget exhausted."""
    config = load_sources_config()
    default_start = config.get("default_start", int(time.time()) - 86400)
    window_seconds = config.get("window_seconds", 3600)

    pipeline = dlt.pipeline(
        pipeline_name="victoriametrics_to_bigquery",
        destination=destination,
        dataset_name=DATASET_NAME,
    )

    run_start = time.time()
    window_num = 0

    while True:
        elapsed = time.time() - run_start
        if elapsed > MAX_RUN_SECONDS:
            logger.info("Time budget exhausted after %d windows (%.0fs)", window_num, elapsed)
            break

        window_num += 1
        rss_mb = resource_mod.getrusage(resource_mod.RUSAGE_SELF).ru_maxrss / 1024
        logger.info("=== Window %d (elapsed %.0fs, peak RSS %.0fMB) ===", window_num, elapsed, rss_mb)

        any_data = False
        for metric_cfg in config["metrics"]:
            source = vm_source(
                metric_cfg=metric_cfg,
                endpoints=config["endpoints"],
                default_start=default_start,
                window_seconds=window_seconds,
            )
            load_info = pipeline.run(source)
            pkg_info = str(load_info)
            if "0 load package" not in pkg_info:
                any_data = True

        if not any_data:
            logger.info("All resources caught up after %d windows", window_num)
            break


if __name__ == "__main__":
    run_pipeline()