#!/usr/bin/env python3
"""
Failed Job Sync - Syncs failed dataflow jobs from Grafana to Bitable
"""

import sys
import requests
from datetime import datetime, date, timedelta
from base import (
    get_bitable_records,
    add_to_bitable,
    FAILED_JOBS_APP_TOKEN,
    FAILED_JOBS_TABLE_ID,
    GRAFANA_TOKEN,
    load_env_from_bashrc
)

# Load env
load_env_from_bashrc()

GRAFANA_LINK_IMPORT = "https://grafana.oauth.prod.aws.tidbcloud.com/d-sources/import-management"
GRAFANA_LINK_DM = "https://grafana.oauth.prod.aws.tidbcloud.com/d/9sVoewCIz/global-control-user-metrics"
GRAFANA_LINK_BACKUP = "https://grafana.oauth.prod.aws.tidbcloud.com"

INTERNAL_TENANTS = (
    "1372813089194211283",
    "1372813089187791289",
    "1372813089189381281",
    "1372813089208161328",
    "1372813089188691290",
    "1372813089209308952",
)


def normalize_failed_reason(value):
    """Keep the full failure text while normalizing null-ish values."""
    if value is None:
        return ""
    return str(value).strip()


def query_grafana(sql):
    """Query Grafana and return rows"""
    if not GRAFANA_TOKEN:
        print("Error: DATASERVICE_GRAFANA_TOKEN not set")
        return []

    payload = {
        "queries": [{
            "refId": "A",
            "datasourceId": 37,
            "rawSql": sql,
            "format": "table"
        }],
        "from": "now-30d",
        "to": "now"
    }

    try:
        resp = requests.post(
            "https://grafana.oauth.prod.aws.tidbcloud.com/api/ds/query",
            headers={"Authorization": f"Bearer {GRAFANA_TOKEN}", "Content-Type": "application/json"},
            json=payload, timeout=120
        )
        data = resp.json()
        frames = data.get("results", {}).get("A", {}).get("frames", [])
        if frames and "data" in frames[0]:
            return frames[0]["data"]["values"]
    except Exception as e:
        print(f"Error querying Grafana: {e}")

    return []


def query_failed_imports(date_str=None):
    """Query failed Import jobs"""
    print("\n=== Querying Failed Import Jobs ===")

    # Filter by date in SQL using string comparison
    if date_str:
        year = int(date_str[:4])
        month = int(date_str[4:6])
        day = int(date_str[6:8])
        target_date = date(year, month, day)
        start_date = target_date.strftime("%Y-%m-%d")
        end_date = (target_date + timedelta(days=1)).strftime("%Y-%m-%d")
        date_filter = f"AND created_at >= '{start_date}' AND created_at < '{end_date}'"
    else:
        # Use string comparison for yesterday
        date_filter = "AND created_at >= DATE_SUB(CURDATE(), INTERVAL 1 DAY)"

    sql = f"""SELECT CAST(cluster_id AS CHAR), CAST(id AS CHAR), message, created_at, CAST(tenant_id AS CHAR), CAST(project_id AS CHAR) FROM imports WHERE status = 'failed' {date_filter} AND tenant_id NOT IN ({','.join(INTERNAL_TENANTS)}) ORDER BY created_at DESC"""

    rows = query_grafana(sql)
    if not rows:
        return []

    # Process results
    results = []
    for i in range(len(rows[0])):
        cluster = str(rows[0][i]).strip() if rows[0][i] else ""
        job_id = str(rows[1][i]) if rows[1][i] else ""
        message = normalize_failed_reason(rows[2][i])
        created_at = rows[3][i] if len(rows) > 3 and rows[3][i] else 0
        tenant_id = str(rows[4][i]) if len(rows) > 4 and rows[4][i] else ""
        project_id = str(rows[5][i]) if len(rows) > 5 and rows[5][i] else ""

        # Determine error type
        if "PreCheck" in message:
            error_type = "PreCheckFailed"
        elif "checksum" in message.lower():
            error_type = "ChecksumMismatch"
        else:
            error_type = "ErrRestoreTable"

        results.append({
            "cluster": cluster,
            "job_id": job_id,
            "error_type": error_type,
            "message": message,
            "created_at": created_at,
            "tenant_id": tenant_id,
            "project_id": project_id,
            "component": "Import"
        })

    print(f"Found {len(results)} failed Import records for {date_str or 'recent'}")
    return results


def query_failed_dm(date_str=None):
    """Query failed DM jobs"""
    print("\n=== Querying Failed DM Jobs ===")

    # Build date filter
    if date_str:
        year = date_str[:4]
        month = date_str[4:6]
        day = date_str[6:8]
        start_date = f"{year}-{month}-{day}"
        next_day = int(day) + 1
        if next_day > 31:
            end_date = f"{int(year)+1}-01-01"
        else:
            end_date = f"{year}-{month}-{next_day:02d}"
        date_filter = f"AND a.created_at >= '{start_date}' AND a.created_at < '{end_date}'"
    else:
        # Use string comparison for yesterday
        date_filter = "AND a.created_at >= DATE_SUB(CURDATE(), INTERVAL 1 DAY)"

    # Query DM jobs with errors - use status = 'Error' filter
    sql = f"""
    SELECT *
    FROM (
        SELECT
            a.id,
            IF(
                a.error_message = ""
                && (JSON_EXTRACT(c.dump_status, '$.errorMsg') IS NULL OR JSON_TYPE(JSON_EXTRACT(c.dump_status, '$.errorMsg')) = 'NULL')
                && (JSON_EXTRACT(c.load_status, '$.errorMsg') IS NULL OR JSON_TYPE(JSON_EXTRACT(c.load_status, '$.errorMsg')) = 'NULL')
                && (JSON_EXTRACT(c.sync_status, '$.errorMsg') IS NULL OR JSON_TYPE(JSON_EXTRACT(c.sync_status, '$.errorMsg')) = 'NULL'),
                'Normal',
                'Error'
            ) AS status,
            concat(a.cluster_id, '') as cluster_id,
            concat(a.project_id, '') as project_id,
            concat(a.tenant_id, '') as tenant_id,
            a.created_at,
            a.error_message,
            c.dump_status,
            c.load_status,
            c.sync_status
        FROM df_jobs a
        JOIN dm_jobs b ON a.id = b.id
        JOIN dm_job_sources c ON a.id = c.job_id
        WHERE
            a.service_code = 'dm'
            AND a.tenant_id NOT IN ({','.join(INTERNAL_TENANTS)})
            {date_filter}
    ) t
    WHERE status = 'Error'
    ORDER BY id DESC
    LIMIT 1000
    """

    rows = query_grafana(sql)
    if not rows:
        return []

    # Filter by date and error status in Python
    results = []
    if date_str:
        year = int(date_str[:4])
        month = int(date_str[4:6])
        day = int(date_str[6:8])
        date(year, month, day)

    # Parse the new query results
    results = []
    for i in range(len(rows[0])):
        job_id = str(rows[0][i]) if rows[0][i] else ""
        status = str(rows[1][i]) if len(rows) > 1 and rows[1][i] else "Normal"
        cluster = str(rows[2][i]).strip() if len(rows) > 2 and rows[2][i] else ""
        project_id = str(rows[3][i]) if len(rows) > 3 and rows[3][i] else ""
        tenant_id = str(rows[4][i]) if len(rows) > 4 and rows[4][i] else ""
        created_at = rows[5][i] if len(rows) > 5 and rows[5][i] else 0
        error_message = normalize_failed_reason(rows[6][i] if len(rows) > 6 else "")
        
        # Build error message from status fields
        message = error_message
        if not message:
            dump_status = rows[7][i] if len(rows) > 7 and rows[7][i] else "{}"
            load_status = rows[8][i] if len(rows) > 8 and rows[8][i] else "{}"
            sync_status = rows[9][i] if len(rows) > 9 and rows[9][i] else "{}"
            message = normalize_failed_reason(f"dump:{dump_status}, load:{load_status}, sync:{sync_status}")
        
        # Include all records with status = 'Error'
        if status == 'Error':
            results.append({
                "cluster": cluster,
                "job_id": job_id,
                "error_type": "DMFailed",
                "message": message,
                "created_at": created_at,
                "tenant_id": tenant_id,
                "project_id": project_id,
                "component": "DM"
            })
    
    print(f"Found {len(results)} failed DM records for {date_str or 'recent'}")
    return results


def query_failed_backups(date_str=None):
    """Query failed Backup jobs."""
    print("\n=== Querying Failed Backup Jobs ===")

    if date_str:
        target_date = date(int(date_str[:4]), int(date_str[4:6]), int(date_str[6:8]))
    else:
        target_date = date.today() - timedelta(days=1)

    sql = f"""
    SELECT
      CONCAT(id, '') AS id,
      created_at,
      updated_at,
      deleted_at,
      CONCAT(tenant_id, '') AS tenant_id,
      CONCAT(project_id, '') AS project_id,
      CONCAT(cluster_id, '') AS cluster_id,
      err_msg,
      CONCAT(backup_schedule_id, '') AS backup_schedule_id,
      targets,
      cluster_version,
      status,
      cluster_config,
      cleaned_up,
      old_backup_name,
      finished_at,
      start_time_stamp,
      backup_size,
      CONCAT(cluster_snapshot_id, '') AS cluster_snapshot_id,
      versioning,
      storage,
      CONCAT(bucket_id, '') AS bucket_id,
      clean_policy,
      is_volume_backup,
      leak_checked,
      pi_tr_setting_snapshot,
      bucket_setting_snapshot,
      table_filter,
      expected_backup_ts,
      is_internal,
      protected,
      extra_info
    FROM backups
    WHERE status IN ('failed')
    ORDER BY created_at DESC
    LIMIT 1000
    """

    rows = query_grafana(sql)
    if not rows:
        return []

    results = []
    start_ts = int(datetime.combine(target_date, datetime.min.time()).timestamp() * 1000)
    end_ts = int(datetime.combine(target_date + timedelta(days=1), datetime.min.time()).timestamp() * 1000)
    internal_tenants = set(INTERNAL_TENANTS)
    for i in range(len(rows[0])):
        job_id = str(rows[0][i]) if rows[0][i] else ""
        created_at = rows[1][i] if len(rows) > 1 and rows[1][i] else 0
        tenant_id = str(rows[4][i]) if len(rows) > 4 and rows[4][i] else ""
        project_id = str(rows[5][i]) if len(rows) > 5 and rows[5][i] else ""
        cluster = str(rows[6][i]).strip() if len(rows) > 6 and rows[6][i] else ""
        message = normalize_failed_reason(rows[7][i] if len(rows) > 7 else "")

        if not (start_ts <= int(created_at) < end_ts):
            continue
        if tenant_id in internal_tenants:
            continue

        results.append({
            "cluster": cluster,
            "job_id": job_id,
            "error_type": "BackupFailed",
            "message": message or "Backup job status is failed with empty err_msg.",
            "created_at": created_at,
            "tenant_id": tenant_id,
            "project_id": project_id,
            "component": "Backup",
        })

    print(f"Found {len(results)} failed Backup records for {date_str or 'recent'}")
    return results


def sync(date_str=None):
    """Main sync function"""
    print("\n" + "=" * 50)
    print("Failed Job Sync" + (f" - {date_str}" if date_str else ""))
    print("=" * 50)

    # Get existing records for deduplication
    existing_keys = get_bitable_records(FAILED_JOBS_APP_TOKEN, FAILED_JOBS_TABLE_ID, ["Job ID"])
    print(f"Existing records in Bitable: {len(existing_keys)}")

    total_added = 0

    # Process Import jobs
    import_jobs = query_failed_imports(date_str)
    for job in import_jobs:
        unique_key = job["job_id"] + "|"
        if unique_key in existing_keys:
            print(f"SKIP: {unique_key} already exists")
            continue

        timestamp = job["created_at"] if job["created_at"] else int(datetime.now().timestamp() * 1000)

        record = {
            "Time": int(timestamp),
            "Component": job["component"],
            "Cluster": job["cluster"],
            "Job ID": job["job_id"],
            "Failed Tasks": job["error_type"],
            "Failed Reason": job["message"],
            "Grafana Link": {"text": "Grafana", "link": GRAFANA_LINK_IMPORT},
            "Tenant ID": job["tenant_id"],
            "Project ID": job["project_id"],
        }

        if add_to_bitable(FAILED_JOBS_APP_TOKEN, FAILED_JOBS_TABLE_ID, record):
            print(f"ADDING: {job['job_id']} - {job['error_type']}")
            total_added += 1
        existing_keys[unique_key] = True

    # Process DM jobs
    dm_jobs = query_failed_dm(date_str)
    for job in dm_jobs:
        unique_key = job["job_id"] + "|"
        if unique_key in existing_keys:
            print(f"SKIP: {unique_key} already exists")
            continue

        timestamp = job["created_at"] if job["created_at"] else int(datetime.now().timestamp() * 1000)

        record = {
            "Time": int(timestamp),
            "Component": job["component"],
            "Cluster": job["cluster"],
            "Job ID": job["job_id"],
            "Failed Tasks": job["error_type"],
            "Failed Reason": job["message"],
            "Grafana Link": {"text": "Grafana", "link": GRAFANA_LINK_DM},
            "Tenant ID": job["tenant_id"],
            "Project ID": job["project_id"],
        }

        if add_to_bitable(FAILED_JOBS_APP_TOKEN, FAILED_JOBS_TABLE_ID, record):
            print(f"ADDING: {job['job_id']} - {job['error_type']}")
            total_added += 1
        existing_keys[unique_key] = True

    # Process Backup jobs
    backup_jobs = query_failed_backups(date_str)
    for job in backup_jobs:
        unique_key = job["job_id"] + "|"
        if unique_key in existing_keys:
            print(f"SKIP: {unique_key} already exists")
            continue

        timestamp = job["created_at"] if job["created_at"] else int(datetime.now().timestamp() * 1000)

        record = {
            "Time": int(timestamp),
            "Component": job["component"],
            "Cluster": job["cluster"],
            "Job ID": job["job_id"],
            "Failed Tasks": job["error_type"],
            "Failed Reason": job["message"],
            "Grafana Link": {"text": "Grafana", "link": GRAFANA_LINK_BACKUP},
            "Tenant ID": job["tenant_id"],
            "Project ID": job["project_id"],
        }

        if add_to_bitable(FAILED_JOBS_APP_TOKEN, FAILED_JOBS_TABLE_ID, record):
            print(f"ADDING: {job['job_id']} - {job['error_type']}")
            total_added += 1
        existing_keys[unique_key] = True

    print("\n=== SUMMARY ===")
    print(f"Failed Import: {len(import_jobs)}")
    print(f"Failed DM: {len(dm_jobs)}")
    print(f"Failed Backup: {len(backup_jobs)}")
    print(f"Added: {total_added}")
    print("=== END ===")

    return {
        "import": len(import_jobs),
        "dm": len(dm_jobs),
        "backup": len(backup_jobs),
        "added": total_added
    }


if __name__ == "__main__":
    date_arg = sys.argv[1] if len(sys.argv) > 1 else None
    sync(date_arg)