"""
Script to:
1. Get all cluster vmaccount and vmproject via Kubernetes API
2. Use this data to access corresponding HTTP API to fetch data
"""

import json
import logging
import os
import sys
import time
from datetime import datetime, timedelta
from typing import List, Optional, Tuple

import boto3
import pymysql
import requests
import yaml
from kubernetes import client, config
from kubernetes.client.rest import ApiException

from delta_sink import SourceInfo, write_tsdb_status_to_deltalake, _delta_table_exists, delete_date_from_deltalake

logger = logging.getLogger(__name__)


def setup_logging(level: str = "INFO") -> None:
    # Millisecond precision helps correlate with other systems.
    logging.basicConfig(
        level=getattr(logging, level.upper(), logging.INFO),
        format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        force=True,
    )

def _date_list(start: str, end: str) -> List[str]:
    """Inclusive date list from start to end (YYYYMMDD)."""
    s = datetime.strptime(start, "%Y%m%d").date()
    e = datetime.strptime(end, "%Y%m%d").date()
    days = (e - s).days
    return [(s + timedelta(days=i)).strftime("%Y%m%d") for i in range(days + 1)]


def load_k8s_config():
    """
    Load Kubernetes configuration
    Try in-cluster config first (when running inside cluster), then fallback to kubeconfig
    """
    try:
        # Try in-cluster config first (for pods running inside cluster)
        config.load_incluster_config()
        logger.info("Using in-cluster Kubernetes configuration")
    except config.ConfigException:
        try:
            # Fallback to kubeconfig (for local development)
            config.load_kube_config()
            logger.info("Using kubeconfig for Kubernetes configuration")
        except config.ConfigException as e:
            logger.error("Failed to load Kubernetes configuration: %s", e)
            sys.exit(1)

def _get_vendor(secret_id: str) -> str:
    """
    Determine cloud vendor for RDS secret.
    Prefer env VENDOR (aws / alicloud); else infer from secret identifier format.
    Matches regional-server: AWS Secrets Manager vs AliCloud KMS.
    """
    env_vendor = (os.environ.get("VENDOR") or "").strip().lower()
    if env_vendor in ("aws", "alicloud"):
        return env_vendor
    if secret_id.startswith("arn:aws:secretsmanager:"):
        return "aws"
    if secret_id.startswith("acs:kms:"):
        return "alicloud"
    logger.warning("Unknown secret id format, assuming aws: %s", secret_id[:80])
    return "aws"


def _get_rds_secret_aws(secret_arn: str) -> Optional[dict]:
    """
    Fetch RDS config JSON from AWS Secrets Manager.
    Same contract as regional-server getRdsConfigWithAWS.
    """
    parts = secret_arn.split(":")
    if len(parts) < 4:
        logger.error("Failed to derive region from ARN: %s", secret_arn)
        return None
    region = parts[3]
    try:
        session = boto3.session.Session()
        sm = session.client("secretsmanager", region_name=region)
        result = sm.get_secret_value(SecretId=secret_arn)
        secret_str = result.get("SecretString")
        if not secret_str:
            logger.error("RDS secret %s has empty SecretString", secret_arn)
            return None
        return json.loads(secret_str)
    except Exception as e:
        logger.error("Failed to get RDS secret from AWS Secrets Manager: %s", e)
        return None


def _get_rds_secret_aliyun(secret_id: str, region: str) -> Optional[dict]:
    """
    Fetch RDS config JSON from AliCloud KMS (GetSecretValue).
    Same contract as regional-server getRdsConfigWithAliCloud.
    Uses default credential chain (RRSA when run with same ServiceAccount as regional-server).
    Secret id format: acs:kms:<region>:<account>:secret/*
    """
    try:
        from alibabacloud_credentials.client import Client as CredentialClient
        from alibabacloud_kms20160120.client import Client
        from alibabacloud_tea_openapi import models as open_api_models
    except ImportError as e:
        logger.error(
            "AliCloud SDK not installed (pip install alibabacloud-credentials alibabacloud-kms20160120): %s",
            e,
        )
        return None
    # Default credential chain: env vars, then OIDC/RRSA (same as regional-server Go SDK).
    # When run with serviceAccountName: regional-server-controller-manager and RRSA enabled,
    # no ALIBABA_CLOUD_ACCESS_KEY_* env vars are required.
    credentials_client = CredentialClient()
    config = open_api_models.Config(
        credential=credentials_client,
        region_id=region,
    )
    client = Client(config)
    try:
        from alibabacloud_kms20160120 import models as kms_models

        request = kms_models.GetSecretValueRequest(secret_name=secret_id)
        response = client.get_secret_value(request)
        body = response.body
        secret_data = (
            getattr(body, "secret_data", None) or getattr(body, "SecretData", None)
            if body
            else None
        )
        if not secret_data:
            logger.error("RDS secret %s has empty SecretData", secret_id[:80])
            return None
        return json.loads(secret_data)
    except Exception as e:
        logger.error("Failed to get RDS secret from AliCloud KMS: %s", e)
        return None


def _rds_config_to_connection(cfg: dict, secret_label: str) -> Optional[pymysql.connections.Connection]:
    """Build MySQL connection from RDS config dict (structs.RdsConfig compatible)."""
    host = cfg.get("host")
    port = int(cfg.get("port", 3306))
    user = cfg.get("username")
    password = cfg.get("password")
    dbname = cfg.get("dbname")
    if not all([host, port, user, password, dbname]):
        logger.error("RDS secret %s missing required fields (host/port/username/password/dbname)", secret_label)
        return None
    try:
        conn = pymysql.connect(
            host=host,
            port=port,
            user=user,
            password=password,
            database=dbname,
            cursorclass=pymysql.cursors.DictCursor,
        )
        return conn
    except Exception as e:
        logger.error("Failed to connect to RDS: %s", e)
        return None


def _get_rds_connection():
    """
    Return a MySQL connection to RDS.

    Logic (aligned with regional-server):
      - Read rdsSecretArn from regional-server ConfigMap.
      - Vendor: env VENDOR or infer from secret id (arn:aws:... -> AWS, acs:kms:... -> AliCloud).
      - AWS: fetch config from AWS Secrets Manager.
      - AliCloud: fetch config from AliCloud KMS GetSecretValue.
      - JSON structure is compatible with structs.RdsConfig.
    """
    try:
        secret_id = _get_rds_secret_from_configmap()
    except Exception as e:
        logger.error("Failed to get RDS secret from ConfigMap: %s", e)
        return None

    vendor = _get_vendor(secret_id)
    cfg = None
    if vendor == "aws":
        parts = secret_id.split(":")
        region = parts[3] if len(parts) >= 4 else ""
        cfg = _get_rds_secret_aws(secret_id)
    elif vendor == "alicloud":
        # acs:kms:<region>:<account>:secret/*
        parts = secret_id.split(":")
        region = parts[2] if len(parts) > 2 else os.environ.get("ALIBABA_CLOUD_REGION_ID", "")
        cfg = _get_rds_secret_aliyun(secret_id, region)
    else:
        logger.error("Unsupported vendor for RDS secret: %s", vendor)
        return None

    if not cfg:
        return None
    conn = _rds_config_to_connection(cfg, secret_id[:80])
    if conn:
        logger.info("Connected to RDS for serverless cluster info (vendor=%s)", vendor)
    return conn


def _get_rds_secret_from_configmap() -> str:
    """
    Read RDS secret ARN from regional-server ConfigMap, to avoid per-region env changes.

    ConfigMap:
      - namespace: observability
      - name: regionalserver
      - key: controllerconfig.yaml

    controllerconfig.yaml structure (YAML):
      serverlessConfig:
        rdsSecretArn: <arn>
    """
    # Reuse in-cluster kube config; falls back to kubeconfig if running locally.
    load_k8s_config()

    v1 = client.CoreV1Api()
    cm = v1.read_namespaced_config_map(name="regionalserver", namespace="observability")
    data = cm.data or {}
    controller_cfg = data.get("controllerconfig.yaml")
    if not controller_cfg:
        raise RuntimeError("controllerconfig.yaml not found in ConfigMap regionalserver")

    cfg = yaml.safe_load(controller_cfg) or {}
    serverless_cfg = cfg.get("serverlessConfig") or {}
    secret_arn = serverless_cfg.get("rdsSecretArn")
    if not secret_arn:
        raise RuntimeError("serverlessConfig.rdsSecretArn not found in controllerconfig.yaml")

    logger.info("Loaded RDS secret ARN from ConfigMap regionalserver")
    return secret_arn


def _get_serverless_accounts_from_rds(conn) -> List[Tuple[str, str, str, str, str]]:
    """
    Get serverless parent cluster vm account/project and tenancy info from RDS.

    This mirrors the logic in Go collector, which queries:
      - table `clusters` for biz_type / tenant / project_id / application_id
      - table `cluster_specs` for vm_account_id / vm_project_id and labels
    """
    accounts: List[Tuple[str, str, str, str, str]] = []
    if conn is None:
        return accounts

    sql = """
SELECT
  c.tenant,
  c.project_id,
  c.application_id,
  c.biz_type,
  s.vm_account_id,
  s.vm_project_id
FROM clusters c
JOIN cluster_specs s ON c.cluster_id = s.cluster_id
WHERE s.vm_account_id IS NOT NULL
  AND (s.vm_project_id IS NULL OR s.vm_project_id = 0)
"""
    try:
        with conn.cursor() as cursor:
            cursor.execute(sql)
            rows = cursor.fetchall()
            for row in rows:
                account_id = str(row.get("vm_account_id", ""))
                # project_id always 0 for serverless parent clusters
                project_id = str(0)
                tenant_id = row.get("tenant") or ""
                application_id = row.get("application_id") or ""
                biztype = row.get("biz_type") or ""

                if account_id and project_id:
                    accounts.append((account_id, project_id, tenant_id, application_id, biztype))

        logger.info("Fetched %d serverless parent clusters from RDS", len(accounts))
    except Exception as e:
        logger.error("Failed to query serverless parent clusters from RDS: %s", e)

    return accounts


# It's a hacker way and not very accurate,
# we will use `/admin/tenants?start=xxx&end=xxx` api to get tenants after https://github.com/VictoriaMetrics/VictoriaMetrics/issues/10295 fixed
def get_cluster_accounts(crd_group: str = "observability.pingcap.com", crd_version: str = "v1", crd_plural: str = "clusters") -> List[Tuple[str, str, str, str, str]]:
    """
    Get all cluster accountID, projectID, tenantID, applicationID and biztype.

    For non-serverless clusters, data is read from Kubernetes CRs.
    For serverless clusters, vm account/project and tenancy info are read from RDS.
    
    Args:
        crd_group: CRD group (default: "observability.pingcap.com")
        crd_version: CRD version (default: "v1")
        crd_plural: CRD plural name (default: "clusters")
    
    Returns:
        List[Tuple[str, str, str, str, str]]: List of [(accountID, projectID, tenantID, applicationID, biztype), ...]
    """
    accounts: List[Tuple[str, str, str, str, str]] = []

    # 1) Non-serverless clusters from Kubernetes CRs.
    try:
        load_k8s_config()

        # Create CustomObjectsApi client for CRD resources
        custom_api = client.CustomObjectsApi()

        group = crd_group
        version = crd_version
        plural = crd_plural

        logger.info("Fetching non-serverless cluster resources from Kubernetes API...")

        # Only query non-serverless shared clusters from CRs.
        label_selector_non_serverless = "observability.pingcap.com/bizType notin (dev-devtier,staging-devtier,prod-devtier),observability.pingcap.com/share-mode=shared"

        clusters = custom_api.list_custom_object_for_all_namespaces(
            group=group,
            version=version,
            resource_plural=plural,
            label_selector=label_selector_non_serverless,
        )

        items = clusters.get("items", [])
        logger.info("Found %d non-serverless clusters from CRs", len(items))

        for cluster in items:
            try:
                spec = cluster.get("spec", {})
                victoria_account = spec.get("victoriaAccount", {})
                account_id = victoria_account.get("accountID", "")
                project_id = victoria_account.get("projectID", 0)

                # Extract tenantID and applicationID from labels
                metadata = cluster.get("metadata", {})
                labels = metadata.get("labels", {})
                tenant_id = labels.get("observability.pingcap.com/tenantID", "")
                application_id = labels.get("observability.pingcap.com/applicationID", "")
                biztype = labels.get("observability.pingcap.com/bizType", "")

                if account_id != "":
                    accounts.append((account_id, project_id, tenant_id, application_id, biztype))
            except (KeyError, AttributeError) as e:
                logger.warning("Failed to parse non-serverless cluster resource: %s", e)
                continue
    except ApiException as e:
        logger.error("Failed to fetch non-serverless cluster resources from Kubernetes API: %s", e)
        logger.error("Status: %s, Reason: %s", getattr(e, "status", None), getattr(e, "reason", None))
        sys.exit(1)
    except Exception as e:
        logger.exception("Error accessing Kubernetes API for non-serverless clusters: %s", e)
        sys.exit(1)

    # 2) Serverless parent clusters from RDS.
    conn = _get_rds_connection()
    try:
        serverless_accounts = _get_serverless_accounts_from_rds(conn)
        accounts.extend(serverless_accounts)
    finally:
        if conn is not None:
            try:
                conn.close()
            except Exception:
                pass

    logger.info("Found %d total account/project combinations", len(accounts))
    return accounts


def fetch_tsdb_status(account_id: str, project_id: str, date: str = None) -> dict:
    """
    Access HTTP API to get tsdb status data
    
    Args:
        account_id: VictoriaMetrics account ID
        project_id: VictoriaMetrics project ID
        date: Date parameter in YYYYMMDD format, if None use current date
    
    Returns:
        dict: JSON data returned by API, or None if failed
    """
    if date is None:
        date = datetime.now().strftime("%Y%m%d")

    # Convert YYYYMMDD to YYYY-MM-DD for API call
    if len(date) == 8 and date.isdigit():
        api_date = f"{date[:4]}-{date[4:6]}-{date[6:8]}"
    else:
        api_date = date
    
    url = f"http://vmselect-vmcluster.observability.svc.cluster.local:8481/select/{account_id}:{project_id}/prometheus/api/v1/status/tsdb?date={api_date}&topN=1000"

    try:
        logger.info("Accessing: %s", url)
        response = requests.get(url, timeout=300)
        response.raise_for_status()
        data = response.json()
        
        # Check if totalSeries is 0, ignore this result if so
        total_series = data.get("data", {}).get("totalSeries", 0)
        if total_series == 0:
            logger.info("Skip account/project %s:%s because totalSeries=0", account_id, project_id)
            return None
        
        return data
    
    except requests.exceptions.RequestException as e:
        logger.warning("Request failed: %s", e)
        return None


def main():
    """Main function"""
    import argparse
    
    parser = argparse.ArgumentParser(description="Get cluster vmaccount and vmproject, and access corresponding HTTP API")
    parser.add_argument(
        "--date",
        type=str,
        default=None,
        help="Date parameter (format: YYYYMMDD), default is current date"
    )
    parser.add_argument(
        "--delta-root",
        type=str,
        default=None,
        help="Delta Lake root directory. If set, write results into Delta tables under this path",
    )
    parser.add_argument(
        "--crd-group",
        type=str,
        default="observability.pingcap.com",
        help="CRD group (default: observability.pingcap.com)"
    )
    parser.add_argument(
        "--crd-version",
        type=str,
        default="v1",
        help="CRD version (default: v1)"
    )
    parser.add_argument(
        "--crd-plural",
        type=str,
        default="clusters",
        help="CRD plural name (default: clusters)"
    )
    parser.add_argument(
        "--log-level",
        type=str,
        default="INFO",
        help="Log level: DEBUG, INFO, WARNING, ERROR (default: INFO)",
    )
    
    args = parser.parse_args()
    setup_logging(args.log_level)
    date_str = args.date or (datetime.now().date() - timedelta(days=1)).strftime("%Y%m%d")
    tableExists = True

    # Determine date range: if delta_root is set and sources table not present, backfill the last 30 days
    date_list = [date_str]
    if args.delta_root:
        sources_path = f"{args.delta_root.rstrip('/')}/sources"
        if not _delta_table_exists(sources_path):
            target_dt = datetime.strptime(date_str, "%Y%m%d").date()
            backfill_start = (target_dt - timedelta(days=29)).strftime("%Y%m%d")
            date_list = _date_list(backfill_start, date_str)
            logger.info(
                "Delta table not found at %s, backfilling last 30 days: %s to %s (%d days)",
                sources_path,
                backfill_start,
                date_str,
                len(date_list),
            )
            tableExists = False
    
    # Get all account/project combinations
    accounts = get_cluster_accounts(
        crd_group=args.crd_group,
        crd_version=args.crd_version,
        crd_plural=args.crd_plural
    )
    
    if not accounts:
        logger.info("No account/project combinations found")
        return

    logger.info("Accounts: %s", accounts)

    # Store all results
    results = {}
    
    # Iterate date list and each account/project
    for current_date in date_list:
        logger.info("=== Processing date %s ===", current_date)

        # Delete existing data for this date before processing
        if tableExists and args.delta_root:
            delete_date_from_deltalake(
                delta_root=args.delta_root,
                date=current_date,
            )

        for account_id, project_id, tenant_id, application_id, biztype in accounts:
            logger.info(
                "Processing: accountID=%s, projectID=%s, tenantID=%s, applicationID=%s, biztype=%s",
                account_id,
                project_id,
                tenant_id,
                application_id,
                biztype,
            )
            data = fetch_tsdb_status(account_id, project_id, current_date)
            
            if data:
                key = f"{current_date}:{account_id}:{project_id}"
                results[key] = {
                    "date": current_date,
                    "accountID": account_id,
                    "projectID": project_id,
                    "tenantID": tenant_id,
                    "applicationID": application_id,
                    "biztype": biztype,
                }

                if args.delta_root:
                    t_write = time.time()
                    write_tsdb_status_to_deltalake(
                        delta_root=args.delta_root,
                        date=current_date,
                        source=SourceInfo(
                            accountID=account_id,
                            projectID=str(project_id),
                            tenantID=tenant_id,
                            applicationID=application_id,
                            biztype=biztype,
                        ),
                        tsdb_status=data,
                    )
                    logger.info(
                        "write_tsdb_status_to_deltalake done (%dms) accountID=%s projectID=%s date=%s",
                        int((time.time() - t_write) * 1000),
                        account_id,
                        project_id,
                        current_date,
                    )
            else:
                logger.warning("Failed to fetch data for date %s", current_date)
    
    # Output results
    import json
    output_data = {
        "date": date_str,
        "dates_processed": date_list,
        "total_accounts": len(accounts),
        "success": len(results),
    }
    
    logger.info("=" * 50)
    logger.info("Results:\n%s", json.dumps(output_data, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()