"""
Metrics fetcher for TiDB Cloud clusters using tuotuo library.
"""

import asyncio
from datetime import datetime
from typing import Dict, List

import pandas as pd
import tuotuo
import yaml
from tuotuo.metrics import QueryParam, QueryRange

from src.database.db import get_cluster_metrics, save_cluster_metrics


async def get_cluster_meta(cluster_id: str):
    """
    Get cluster metadata from clinic.

    Args:
        cluster_id: The cluster ID to query

    Returns:
        Cluster metadata object
    """
    client = tuotuo.client("clinic")
    clusters = [
        cluster async for cluster in client.list_clusters(cluster_id=cluster_id)
    ]
    if len(clusters) == 1:
        return clusters[0]
    else:
        raise ValueError(f"There are {len(clusters)} clusters for ID {cluster_id}.")


def get_cluster_meta_sync(cluster_id: str):
    """
    Synchronous wrapper for get_cluster_meta.
    """
    return asyncio.run(get_cluster_meta(cluster_id))


async def fetch_diagram_data(
    cluster_id: str,
    diagram_queries: List[Dict],
    start: datetime,
    end: datetime,
    clinic_org_id: str,
    region_name: str,
    provider: str,
) -> pd.DataFrame:
    """
    Fetch and join multiple queries for a single diagram using concurrent queries.

    Args:
        cluster_id: The cluster ID
        diagram_queries: List of query configs, each with 'name', 'expr', 'unit'
        start: Start time for the query range
        end: End time for the query range
        clinic_org_id: Clinic organization ID
        region_name: Region name
        provider: Cloud provider

    Returns:
        DataFrame with all metrics joined on the time index
    """
    if not diagram_queries:
        return pd.DataFrame()

    query_range = QueryRange(start, end, "15s")

    # Create all query params
    query_params = []
    for query in diagram_queries:
        # Extract query labels from the query name
        # If name has ":by_{label}", use that label
        # Otherwise use "name" as default (from label_replace)
        query_name = query["name"]
        if ":by_" in query_name:
            # Extract the label name after ":by_"
            label = query_name.split(":by_")[1]
            query_labels = [label]
        else:
            # Use "name" as default label (set by label_replace in expr)
            query_labels = ["name"]

        query_params.append(
            QueryParam(
                time_arg=query_range,
                query_template=query["expr"],
                cluster_id=cluster_id,
                region_name=region_name,
                clinic_org_id=clinic_org_id,
                provider=provider,
                tidb_cluster=cluster_id,
                query_name=query_name,
                query_labels=query_labels,
            )
        )

    # Use tuotuo client's concurrent query capability
    client = tuotuo.client("metrics")
    async with client:
        # Pass all queries at once - client will handle concurrency
        results = await client.query(*query_params)

        # If only one query, results is a single MetricsData, wrap it in a list
        if not isinstance(results, list):
            results = [results]

        # Convert results to dataframes
        valid_dfs = []
        for result in results:
            if result.exception is None and result.frame is not None:
                df = result.frame
                if not df.empty:
                    valid_dfs.append(df)
            else:
                if result.exception:
                    print(
                        f"Error fetching metric {result.query.query_name}: "
                        f"{result.exception}"
                    )

        if not valid_dfs:
            return pd.DataFrame()

        # Join all dataframes on the time index (ds)
        result_df = valid_dfs[0]
        for df in valid_dfs[1:]:
            result_df = result_df.join(df, how="outer", rsuffix="_dup")

        return result_df


def load_promql_config(config_path: str = "config/promqls.yaml") -> Dict:
    """
    Load promql configuration from yaml file.

    Args:
        config_path: Path to the promql yaml config file

    Returns:
        Dict with diagram names as keys and list of query configs as values
    """
    with open(config_path, "r") as f:
        return yaml.safe_load(f)


async def fetch_all_diagrams_for_cluster(
    cluster_id: str,
    promql_config: Dict,
) -> Dict[str, Dict]:
    """
    Fetch all diagrams for a single cluster.

    Args:
        cluster_id: The cluster ID
        promql_config: The full promql configuration dict

    Returns:
        Dict mapping diagram names to their data dict with:
        - 'dataframe': pd.DataFrame
        - 'queries': list of query configs (with name, expr, unit)
    """
    # Get cluster metadata
    cluster_meta = await get_cluster_meta(cluster_id)

    clinic_org_id = cluster_meta.clinic_org_id
    region = cluster_meta.region_name
    provider = cluster_meta.provider
    start = datetime.fromtimestamp(cluster_meta.created_at / 1000)
    end = (
        datetime.fromtimestamp(cluster_meta.deleted_at / 1000)
        if cluster_meta.deleted_at > 0
        else datetime.now()
    )

    # Fetch all diagrams
    diagrams = {}
    for diagram_name, queries in promql_config.items():
        if isinstance(queries, list):
            df = await fetch_diagram_data(
                cluster_id, queries, start, end, clinic_org_id, region, provider
            )
            diagrams[diagram_name] = {"dataframe": df, "queries": queries}

    return diagrams


def fetch_all_diagrams_sync(
    cluster_id: str, config_path: str = "config/promqls.yaml"
) -> Dict[str, Dict]:
    """
    Synchronous wrapper for fetching all diagrams.

    Args:
        cluster_id: The cluster ID
        config_path: Path to the promql yaml config file

    Returns:
        Dict mapping diagram names to their data dict with:
        - 'dataframe': pd.DataFrame
        - 'queries': list of query configs (with name, expr, unit)
    """
    promql_config = load_promql_config(config_path)
    return asyncio.run(fetch_all_diagrams_for_cluster(cluster_id, promql_config))


def fetch_all_diagrams_with_cache(
    task_id: int,
    cluster_id: str,
    config_path: str = "config/promqls.yaml",
) -> Dict[str, Dict]:
    """
    Fetch diagrams with database caching.

    Logic:
    1. Try to get metrics from database cache
    2. If cache exists, return cached data
    3. If cache doesn't exist, fetch from Clinic API,
       save to database, and return

    Args:
        task_id: Task ID (for cache key)
        cluster_id: Cluster ID
        config_path: Path to the promql yaml config file

    Returns:
        Dict mapping diagram names to their data dict with:
        - 'dataframe': pd.DataFrame
        - 'queries': list of query configs (with name, expr, unit)
    """

    # Try to get from cache
    cached_diagrams = get_cluster_metrics(task_id, cluster_id)

    if cached_diagrams is not None:
        print(f"Using cached metrics for task {task_id}, cluster {cluster_id}")
        return cached_diagrams

    # Cache miss - fetch from Clinic API
    print(f"Fetching metrics from Clinic API for cluster {cluster_id}...")
    diagrams = fetch_all_diagrams_sync(cluster_id, config_path)

    # Save to cache
    saved_id = save_cluster_metrics(task_id, cluster_id, diagrams)
    if saved_id:
        print(f"Saved metrics to database (id={saved_id})")
    else:
        print("Warning: Failed to save metrics to database")

    return diagrams