"""
This script fetches spans with duration > 10s and status code 200 for the last 7 days from Datadog API.

Then it collects unique URLs from the spans and makes requests to each unique URL with headers for Insights employee.

The goal is to pre-warm Redis cache for the employees.
"""

import logging
import os
import time
from collections import Counter

import requests
from datadog_api_client import ApiClient, Configuration
from datadog_api_client.v2.api.spans_api import SpansApi
from datadog_api_client.v2.model.spans_list_request import SpansListRequest
from datadog_api_client.v2.model.spans_list_request_attributes import (
    SpansListRequestAttributes,
)
from datadog_api_client.v2.model.spans_list_request_data import SpansListRequestData
from datadog_api_client.v2.model.spans_list_request_page import SpansListRequestPage
from datadog_api_client.v2.model.spans_list_request_type import SpansListRequestType
from datadog_api_client.v2.model.spans_query_filter import SpansQueryFilter
from datadog_api_client.v2.model.spans_query_options import SpansQueryOptions
from owsrequest.constants import headers as owsrequest_headers

logging.basicConfig(level=logging.INFO)

INSIGHTS_EMPLOYEE_REQUEST_HEADERS = {
    owsrequest_headers.ORCHARD_PROFILE_ID: "154",  # todo: replace with a dedicated profile id just in case
    owsrequest_headers.ORCHARD_PROFILE_TYPE: "InsightsProfile",
}

DD_API_CLIENT_CONFIGURATION = Configuration()
DD_API_CLIENT_CONFIGURATION.api_key["apiKeyAuth"] = os.getenv("DATADOG_API_KEY")
DD_API_CLIENT_CONFIGURATION.api_key["appKeyAuth"] = os.getenv("DATADOG_APP_KEY")

body = SpansListRequest(
    data=SpansListRequestData(
        attributes=SpansListRequestAttributes(
            filter=SpansQueryFilter(
                _from="now-7d",
                query="service:ows-analytics AND @duration:>10s AND @http.status_code:200 AND @http.method:GET AND @env:prod",
                to="now",
            ),
            options=SpansQueryOptions(
                timezone="GMT",
            ),
            page=SpansListRequestPage(
                limit=1000,
            ),
        ),
        type=SpansListRequestType.SEARCH_REQUEST,
    ),
)

with ApiClient(DD_API_CLIENT_CONFIGURATION) as api_client:
    api_instance = SpansApi(api_client)
    logging.info(
        "Fetching spans with duration > 10s and status code 200 for the last 7 days..."
    )
    response = api_instance.list_spans(body=body)

    all_requests = []
    unique_users = set()
    unique_urls = set()
    for query in response.data:
        duration_s = int(query["attributes"]["custom"]["duration"] / 1000000000)
        url = query["attributes"]["custom"]["http"]["url"]
        identity_id = (
            query["attributes"]["custom"].get("request_context", {}).get("identity_id")
        )
        if identity_id:
            unique_users.add(identity_id)
            unique_urls.add(url)
            all_requests.append(url)

    all_requests_counter = Counter(all_requests)
    logging.info(f"Grouped requests: {all_requests_counter}")
    logging.info(f"Unique users: {unique_users}")

sum_duration = 0
for url in unique_urls:
    start_time = time.time()
    response = requests.get(url, headers=INSIGHTS_EMPLOYEE_REQUEST_HEADERS)
    end_time = time.time()
    duration = end_time - start_time
    sum_duration += duration
    logging.info(
        f"URL: {url}, Status code: {response.status_code}, Duration: {duration:.2f} seconds"
    )

logging.info(
    f"Total duration: {sum_duration / 60:.2f} minutes prewarming Redis cache for {len(unique_urls)} unique URLs"
)