#!/usr/bin/env python3
import argparse
import json
from pathlib import Path
from typing import Any, Dict

import psycopg2
import yaml


def load_yaml(path: Path) -> Dict[str, Any]:
    return yaml.safe_load(path.read_text(encoding="utf-8"))


def parse_pg_dsn(dsn: str) -> str:
    parts = dsn.split()
    kept = []
    timezone = None
    for part in parts:
        if "=" not in part:
            kept.append(part)
            continue
        k, v = part.split("=", 1)
        if k.lower() == "timezone":
            timezone = v
            continue
        kept.append(part)
    if timezone:
        kept.append(f"options='-c timezone={timezone}'")
    return " ".join(kept)


def main() -> int:
    parser = argparse.ArgumentParser(description="日志管理 E2E 测试数据准备脚本")
    parser.add_argument(
        "--config",
        default="d:/work/ops/logs/etc/logs_dev.yaml",
        help="logs 配置文件路径",
    )
    parser.add_argument("--run-id", required=True, help="本次测试 run id")
    parser.add_argument("--cleanup-only", action="store_true", help="仅清理历史测试数据")
    args = parser.parse_args()

    cfg = load_yaml(Path(args.config))
    dsn = parse_pg_dsn(cfg["Databases"]["Source"][0])
    run_id = args.run_id
    marker = f"%[E2E:{run_id}]%"

    summary: Dict[str, Any] = {"run_id": run_id, "cleanup": {}, "seed": {}}
    with psycopg2.connect(dsn) as conn:
        with conn.cursor() as cur:
            cur.execute("DELETE FROM logs_alert_outbox WHERE id IN (SELECT id FROM logs_alert_outbox ORDER BY id DESC LIMIT 0)")
            cur.execute("DELETE FROM logs_syslog_rules WHERE name LIKE %s", (marker,))
            summary["cleanup"]["logs_syslog_rules"] = cur.rowcount
            cur.execute("DELETE FROM logs_trap_rules WHERE name LIKE %s", (marker,))
            summary["cleanup"]["logs_trap_rules"] = cur.rowcount
            cur.execute("DELETE FROM logs_trap_dictionary WHERE title LIKE %s", (marker,))
            summary["cleanup"]["logs_trap_dictionary"] = cur.rowcount
            cur.execute("DELETE FROM logs_trap_shields WHERE name LIKE %s", (marker,))
            summary["cleanup"]["logs_trap_shields"] = cur.rowcount
            cur.execute(
                "DELETE FROM logs_resource_event_dedup WHERE event_id LIKE %s",
                (f"e2e-{run_id}-%",),
            )
            summary["cleanup"]["logs_resource_event_dedup"] = cur.rowcount
            if not args.cleanup_only:
                cur.execute(
                    """
                    INSERT INTO logs_resource_mappings(
                        resource_type, resource_id, resource_name, ips_json, hostnames_json, labels_json,
                        version, is_deleted, last_event_id, event_time, created_at, updated_at
                    ) VALUES (%s,%s,%s,%s,%s,%s,%s,false,%s,now(),now(),now())
                    ON CONFLICT (resource_type, resource_id)
                    DO UPDATE SET
                        resource_name=EXCLUDED.resource_name,
                        ips_json=EXCLUDED.ips_json,
                        hostnames_json=EXCLUDED.hostnames_json,
                        labels_json=EXCLUDED.labels_json,
                        version=EXCLUDED.version,
                        is_deleted=false,
                        last_event_id=EXCLUDED.last_event_id,
                        event_time=now(),
                        updated_at=now()
                    """,
                    (
                        "server",
                        f"seed-{run_id}",
                        f"E2E-Seed-{run_id}",
                        json.dumps(["127.0.0.1"]),
                        json.dumps([f"e2e-host-{run_id}"]),
                        json.dumps({"source": "prepare_script"}),
                        1,
                        f"e2e-{run_id}-seed",
                    ),
                )
                summary["seed"]["resource_mapping"] = {"resource_type": "server", "resource_id": f"seed-{run_id}"}
        conn.commit()

    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())