from typing import Any, Optional
import yaml
import os

import time
import dlt
from dlt.sources.rest_api import (
    RESTAPIConfig,
    rest_api_resources,
)

from dlt.sources.helpers.rest_client.paginators import SinglePagePaginator

# dlt.config["runtime.log_level"] = "DEBUG"
dlt.config["destination.bigquery.location"] = "us-west2"

# Prometheus API endpoint and sample query
def get_prometheus_config(
    query: str,
    step: int,
    base_url: str,
    resource_name: str,
    start: int = None,
    end: int = None
) -> RESTAPIConfig:
    endpoint_params = {
        "start": start,
        "end": end,
        "query": query,
        "step": step,
    }

    config = {
        "client": {
            "base_url": base_url,
            "paginator": SinglePagePaginator(),
        },
        "resources": [
            {
                "name": resource_name,
                "endpoint": {
                    "path": "query_range",
                    "method": "POST",
                    "params": endpoint_params,
                },
            }
        ],
    }
    
    return config

@dlt.resource(name="prometheus_query_range")
def get_prometheus_resource(query, step, base_url, resource_name, start=None, end=None):
    state = dlt.current.resource_state()
    # Use the last_end from state or set a default value
    start = state.setdefault("last_end", 1740758400)
    end = start + 3600

    print(f"Querying Prometheus from {start} to {end} with step {step} for resource {resource_name}")
    
    current_time = int(time.time())
    latest_end = current_time - (current_time % 3600)  # Align to the nearest hour
    
    if end > latest_end:
        # If the end time exceeds the latest available data, skip for this resource
        yield from []
        
    else:
        config = get_prometheus_config(query, step, base_url, resource_name, start=start, end=end - 1)
        
        for resource in rest_api_resources(config):
            for item in resource:
                if hasattr(item, 'data'):
                    data = getattr(item, 'data')
                else:
                    data = item

                if not isinstance(data, list):
                    data = [data]
                for subitem in data:
                    metric = subitem.get('metric', {})
                    for ts, value in subitem.get('values', []):
                        row = dict(metric)
                        row['metric_timestamp'] = int(ts)
                        row['metric_value'] = float(value)
                        yield row
        
        # Update the state with the new start and end times
        state["last_end"] = end

@dlt.source()
def get_group_source(group: dict) -> Any:
    for query in group["queries"]:
        for url_info in group["base_urls"]:
            base_url = url_info["url"]
            url_name = url_info["name"]
            resource_name = f"prometheus_{group['name']}_{query['name']}_{url_name}"
            _r = get_prometheus_resource(
                query=query["query"],
                step=query.get("step", 15),
                base_url=base_url,
                resource_name=resource_name,
            )

            yield _r.with_name(resource_name)


def load_prometheus() -> None:
    pipeline = dlt.pipeline(
        pipeline_name="prometheus_to_bigquery_mar",
        destination='bigquery',
        dataset_name="raw_prometheus_metrics_mar"
    )

    # Load sources config from YAML
    config_path = os.path.join(os.path.dirname(__file__), "sources.yml")
    with open(config_path, "r") as f:
        sources_config = yaml.safe_load(f)
        
    sources = []
    for group in sources_config["groups"]:
        group_source = get_group_source(group)
        sources.append(group_source)

    # For each group source, we rate limit the requests (if a backfill is ongoing)
    for group_source in sources:
        load_info = pipeline.run(group_source),
        # print(load_info)


if __name__ == "__main__":
    # interval 1746853200 may be unprocessed
    remaining_ts = int(time.time()) - 1747195170
    while remaining_ts > 0:
        load_prometheus()
        time.sleep(5)
        remaining_ts -= 3600