#!/usr/bin/env python3
"""
Batch convert Grafana dashboard JSON files to Kubernetes ConfigMaps.

This script:
1. Reads JSON dashboard files from the origin folder
2. Converts __inputs to templating variables
3. Generates Kubernetes ConfigMap YAML files for Prometheus Operator

Credit: Claude Sonnet 4.5 
"""

import json
import os
import sys
from pathlib import Path
from typing import Any, Dict, List, Union


def replace_datasource_references(obj: Any, old_ref: str, new_ref: str) -> Any:
    """
    Recursively replace datasource references in a nested object.
    
    Args:
        obj: Dictionary, list, or primitive value
        old_ref: Old datasource reference (e.g., "${DS_TEST-CLUSTER}")
        new_ref: New datasource reference (e.g., "${datasource}")
        
    Returns:
        Modified object
    """
    if isinstance(obj, dict):
        return {k: replace_datasource_references(v, old_ref, new_ref) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [replace_datasource_references(item, old_ref, new_ref) for item in obj]
    elif isinstance(obj, str):
        return obj.replace(old_ref, new_ref)
    else:
        return obj


def convert_inputs_to_templating(dashboard: Dict[str, Any], datasource_name: str = "datasource", datasource_value: str = "default") -> Dict[str, Any]:
    """
    Convert __inputs to templating variables in a Grafana dashboard.
    Also replaces all datasource references in the dashboard.
    
    Args:
        dashboard: Dictionary containing the dashboard JSON
        datasource_name: Name to use for the datasource variable (default: "datasource")
        datasource_value: Default value for the datasource variable (default: "default")
        
    Returns:
        Modified dashboard dictionary
    """
    if "__inputs" not in dashboard or not dashboard["__inputs"]:
        return dashboard
    
    # Ensure templating structure exists
    if "templating" not in dashboard:
        dashboard["templating"] = {"list": []}
    elif "list" not in dashboard["templating"]:
        dashboard["templating"]["list"] = []
    
    # Replace title Test-Cluster-TiDB-Summary => TiDB / TiDB-Summary
    dashboard["title"] = dashboard["title"].replace("Test-Cluster-", "TiDB / ")
    
    # Process each input
    for input_var in dashboard["__inputs"]:
        if input_var.get("type") == "datasource":
            old_var_name = input_var["name"]
            
            # Create a datasource templating variable
            templating_var = {
                "name": datasource_name,
                "type": "datasource",
                "label": datasource_name,
                "description": input_var.get("description", ""),
                "query": input_var.get("pluginId", "prometheus"),
                "current": {
                    "value": datasource_value,
                    "text": datasource_value
                },
                "hide": 0,
                "includeAll": False,
                "multi": False,
                "options": [],
                "refresh": 1,
                "regex": "",
                "skipUrlSync": False
            }
            
            # Insert at the beginning of the list so datasource vars come first
            dashboard["templating"]["list"].insert(0, templating_var)
            
            # Replace all references to the old datasource variable with the new one
            old_ref = f"${{{old_var_name}}}"
            new_ref = f"${{{datasource_name}}}"
            dashboard = replace_datasource_references(dashboard, old_ref, new_ref)
    
    return dashboard


def sanitize_name(filename: str) -> str:
    """
    Convert filename to a valid Kubernetes resource name.
    
    Args:
        filename: Original filename
        
    Returns:
        Sanitized name suitable for Kubernetes
    """
    # Remove .json extension
    name = filename.replace('.json', '')
    # Convert to lowercase and replace underscores with hyphens
    name = name.lower().replace('_', '-')
    return name


def create_configmap(dashboard_json: str, dashboard_name: str, namespace: str = "tidb-monitoring") -> Dict[str, Any]:
    """
    Create a Kubernetes ConfigMap YAML structure.
    
    Args:
        dashboard_json: Dashboard JSON as a string
        dashboard_name: Name of the dashboard
        namespace: Kubernetes namespace
        
    Returns:
        Dictionary representing the ConfigMap
    """
    labels = {
        "grafana_dashboard": "1",
        "app.kubernetes.io/name": "grafana",
        "app.kubernetes.io/component": "dashboard",
        "app.kubernetes.io/managed-by": "batch-converter"
    }
    
    configmap = {
        "apiVersion": "v1",
        "kind": "ConfigMap",
        "metadata": {
            "labels": labels,
            "name": dashboard_name,
            "namespace": namespace
        },
        "data": {
            f"{dashboard_name}.json": dashboard_json
        }
    }
    return configmap


def split_dashboard_by_rows(dashboard: Dict[str, Any], max_size: int = 1048576) -> List[Dict[str, Any]]:
    """
    Split a large dashboard into multiple parts by row groups if it exceeds max_size.
    
    Args:
        dashboard: Dashboard dictionary
        max_size: Maximum size in bytes for a single dashboard part
        
    Returns:
        List of dashboard parts (each is a full dashboard with subset of rows)
    """
    dashboard_json = json.dumps(dashboard, separators=(',', ':'))
    
    # If dashboard fits in one ConfigMap, return as-is
    if len(dashboard_json.encode('utf-8')) <= max_size:
        return [dashboard]
    
    # Need to split by rows
    if "panels" not in dashboard or not dashboard["panels"]:
        # No panels to split, return as-is and let it fail
        return [dashboard]
    
    parts = []
    base_dashboard = {k: v for k, v in dashboard.items() if k != "panels"}
    
    # Split panels into row groups
    current_panels = []
    
    for panel in dashboard["panels"]:
        # Create a test dashboard with current panels + new panel
        test_dashboard = base_dashboard.copy()
        test_dashboard["panels"] = current_panels + [panel]
        test_json = json.dumps(test_dashboard, separators=(',', ':'))
        
        if len(test_json.encode('utf-8')) > max_size and current_panels:
            # Current panels would exceed limit, save current part and start new
            part = base_dashboard.copy()
            part["panels"] = current_panels
            parts.append(part)
            current_panels = [panel]
        else:
            current_panels.append(panel)
    
    # Add the remaining panels as the last part
    if current_panels:
        part = base_dashboard.copy()
        part["panels"] = current_panels
        parts.append(part)
    
    if parts:
        for (idx, part) in enumerate(parts):
            part["title"] = part["title"] + f" (Partition {idx+1})"
            part["uid"] = part["uid"] + f"_{idx}"

    return parts if parts else [dashboard]


def process_dashboards(origin_dir: str, output_dir: str, namespace: str = "tidb-monitoring", 
                       datasource_name: str = "datasource", datasource_value: str = "default", tags: Union[List[str], None] = None) -> None:
    """
    Process all dashboard JSON files and create ConfigMap YAML files.
    
    Args:
        origin_dir: Directory containing original JSON files
        output_dir: Directory to write ConfigMap YAML files
        namespace: Kubernetes namespace for ConfigMaps
        datasource_name: Name to use for datasource variable
        datasource_value: Default value for datasource variable
        tags: List of tags to add to dashboard JSON
    """
    origin_path = Path(origin_dir)
    output_path = Path(output_dir)
    
    # Create output directory if it doesn't exist
    output_path.mkdir(parents=True, exist_ok=True)
    
    # Process each JSON file
    json_files = list(origin_path.glob("*.json"))
    
    if not json_files:
        print(f"No JSON files found in {origin_dir}")
        return
    
    print(f"Found {len(json_files)} dashboard files to process")
    
    for json_file in json_files:
        try:
            print(f"Processing {json_file.name}...", end=" ")
            
            # Read the dashboard JSON
            with open(json_file, 'r', encoding='utf-8') as f:
                dashboard = json.load(f)
            
            # Convert __inputs to templating
            dashboard = convert_inputs_to_templating(dashboard, datasource_name, datasource_value)
            
            # Add tags to dashboard JSON if provided
            if tags:
                if "tags" not in dashboard:
                    dashboard["tags"] = []
                # Add tags that don't already exist
                for tag in tags:
                    if tag not in dashboard["tags"]:
                        dashboard["tags"].append(tag)
            
            for template in dashboard["templating"]["list"]:
                if template["label"] == "tidb_cluster":
                    template["hide"] = 0
            
            # Split dashboard if too large
            dashboard_parts = split_dashboard_by_rows(dashboard)
            
            # Generate ConfigMap name
            base_configmap_name = sanitize_name(json_file.name)
            
            # Process each part
            for i, dashboard_part in enumerate(dashboard_parts):
                # Convert dashboard back to JSON string (compact format)
                dashboard_json = json.dumps(dashboard_part, separators=(',', ':'))
                
                # Add part suffix if split into multiple parts
                if len(dashboard_parts) > 1:
                    configmap_name = f"{base_configmap_name}-part{i+1}"
                    print(f"[Part {i+1}/{len(dashboard_parts)}]", end=" ")
                else:
                    configmap_name = base_configmap_name
                
                # Create ConfigMap structure
                configmap = create_configmap(dashboard_json, configmap_name, namespace)
                
                # Write ConfigMap YAML file
                output_file = output_path / f"{configmap_name}.yml"
                with open(output_file, 'w', encoding='utf-8') as f:
                    # Use custom YAML formatting for the dashboard JSON
                    # to match the format in the example
                    labels = configmap['metadata']['labels']
                    yaml_lines = [
                        f"apiVersion: {configmap['apiVersion']}\n",
                        f"kind: {configmap['kind']}\n",
                        "metadata:\n",
                        "  labels:\n",
                    ]
                    
                    # Write all labels dynamically
                    for label_key, label_value in labels.items():
                        yaml_lines.append(f"    {label_key}: \"{label_value}\"\n")
                    
                    yaml_lines.extend([
                        f"  name: {configmap['metadata']['name']}\n",
                        f"  namespace: {configmap['metadata']['namespace']}\n",
                        "data:\n",
                        f"  \"{configmap_name}.json\": |\n",
                        f"    {dashboard_json}\n"
                    ])
                    f.writelines(yaml_lines)
            
            if len(dashboard_parts) == 1:
                print(f"✓ Created {base_configmap_name}.yml")
            else:
                print(f"✓ Created {len(dashboard_parts)} parts")
            
        except Exception as e:
            print(f"✗ Error: {e}")
            continue
    
    print(f"\nProcessing complete! Output files in: {output_dir}")


def main() -> None:
    """Main function to handle CLI usage."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Convert Grafana dashboards to Kubernetes ConfigMaps"
    )
    parser.add_argument(
        "--origin-dir",
        default="./origin",
        help="Directory containing original JSON dashboard files (default: ./origin)"
    )
    parser.add_argument(
        "--output-dir",
        default="./configmaps",
        help="Directory to write ConfigMap YAML files (default: ./configmaps)"
    )
    parser.add_argument(
        "--namespace",
        default="tidb-monitoring",
        help="Kubernetes namespace for ConfigMaps (default: tidb-monitoring)"
    )
    parser.add_argument(
        "--datasource-name",
        default="datasource",
        help="Name for the datasource variable in dashboards (default: datasource)"
    )
    parser.add_argument(
        "--datasource-value",
        default="default",
        help="Default value for the datasource variable (default: default)"
    )
    parser.add_argument(
        "--tag",
        action="append",
        dest="tags",
        help="Add tags to dashboard JSON (can be used multiple times)"
    )
    
    args = parser.parse_args()
    
    # Convert to absolute paths
    origin_dir = os.path.abspath(args.origin_dir)
    output_dir = os.path.abspath(args.output_dir)
    
    if not os.path.exists(origin_dir):
        print(f"Error: Origin directory '{origin_dir}' does not exist")
        sys.exit(1)
    
    process_dashboards(origin_dir, output_dir, args.namespace, 
                      args.datasource_name, args.datasource_value, args.tags)


if __name__ == "__main__":
    main()