"""
Backfill pipeline: loads historical data (~90 days) into the same BigQuery tables
as the main pipeline, but with its own state tracking (pipeline_name="vm_backfill").

Processes one metric at a time to stay within memory limits on e2-medium.
Self-terminates once it reaches the main pipeline's start timestamp.
"""

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_backfill")

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

BACKFILL_TARGET = 1773700800  # main pipeline's default_start (2026-03-14 ~22:40 UTC)


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,
):
    @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

        if end > BACKFILL_TARGET:
            logger.info(
                "[%s] Backfill complete (reached target %d), skipping",
                resource_name, BACKFILL_TARGET,
            )
            return

        logger.info("[%s] backfill 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:
    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 = 18 * 60  # stay within 20-min cron interval


def run_backfill() -> None:
    """Loops through windows until backfill target reached or time budget exhausted."""
    config = load_sources_config()
    default_start = BACKFILL_TARGET - (90 * 86400)
    window_seconds = config.get("window_seconds", 3600)

    pipeline = dlt.pipeline(
        pipeline_name="vm_backfill",
        destination="bigquery",
        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("=== Backfill 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("Backfill complete after %d windows", window_num)
            break


if __name__ == "__main__":
    run_backfill()