#!/usr/bin/env python3
"""Register the [Agent] event schema in your Amplitude project's data catalog.

This CLI populates your project's data catalog with all [Agent] event types,
property names, types, descriptions, and required flags via the Taxonomy API.
It is idempotent and safe to re-run.

Usage:
    amplitude-ai-register-catalog --api-key YOUR_KEY --secret-key YOUR_SECRET
    amplitude-ai-register-catalog --dry-run   # preview without API calls

Requirements:
    - Amplitude Enterprise plan (Taxonomy API access)
    - Project API key and Secret key from Settings > Projects
"""

from __future__ import annotations

import argparse
import base64
import json
import logging
import sys
import urllib.error
import urllib.request
from importlib import resources
from pathlib import Path
from typing import Any
from urllib.parse import quote

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)-8s %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("amplitude-ai-register-catalog")

BASE_URL = "https://amplitude.com/api/2"
CATEGORY = "Agent Analytics"


def _make_auth_header(api_key: str, secret_key: str) -> str:
    token = base64.b64encode(f"{api_key}:{secret_key}".encode()).decode()
    return f"Basic {token}"


def _request(
    method: str,
    path: str,
    data: dict[str, Any] | None,
    auth_header: str,
    dry_run: bool,
    stats: dict[str, int],
) -> dict[str, Any] | None:
    url = f"{BASE_URL}/{path}"

    if dry_run:
        logger.info("[DRY RUN] %s %s  data=%s", method, url, data)
        return {"success": True}

    encoded = urllib.parse.urlencode(data).encode() if data else None
    req = urllib.request.Request(url, data=encoded, method=method)
    req.add_header("Authorization", auth_header)
    if encoded:
        req.add_header("Content-Type", "application/x-www-form-urlencoded")

    try:
        with urllib.request.urlopen(req) as resp:
            body: dict[str, Any] = json.loads(resp.read().decode())
            return body
    except urllib.error.HTTPError as e:
        body_text = e.read().decode() if e.fp else ""
        if e.code == 409:
            return {"already_exists": True}
        logger.error("%s %s -> HTTP %d: %s", method, url, e.code, body_text)
        stats["errors"] += 1
        return None
    except Exception:
        logger.exception("Request failed: %s %s", method, url)
        stats["errors"] += 1
        return None


def _ensure_category(auth_header: str, dry_run: bool, stats: dict[str, int]) -> None:
    _request("POST", "taxonomy/category", {"name": CATEGORY}, auth_header, dry_run, stats)


def _register_event(
    event_type: str,
    description: str,
    auth_header: str,
    dry_run: bool,
    stats: dict[str, int],
) -> None:
    result = _request(
        "POST",
        "taxonomy/event",
        {"event_type": event_type, "description": description, "category": CATEGORY},
        auth_header,
        dry_run,
        stats,
    )
    if result and result.get("already_exists"):
        encoded_name = quote(event_type)
        update_result = _request(
            "PUT",
            f"taxonomy/event/{encoded_name}",
            {"description": description, "category": CATEGORY},
            auth_header,
            dry_run,
            stats,
        )
        if update_result:
            stats["events_updated"] += 1
    elif result:
        stats["events_created"] += 1


def _register_property(
    event_type: str,
    prop: dict[str, Any],
    auth_header: str,
    dry_run: bool,
    stats: dict[str, int],
) -> None:
    data: dict[str, Any] = {
        "event_type": event_type,
        "event_property": prop["name"],
        "description": prop.get("description", ""),
        "type": prop.get("type", "string"),
    }
    if prop.get("is_required"):
        data["is_required"] = "true"
    if prop.get("is_array_type"):
        data["is_array_type"] = "true"

    result = _request("POST", "taxonomy/event-property", data, auth_header, dry_run, stats)
    if result and result.get("already_exists"):
        update_data = {k: v for k, v in data.items() if k != "event_type"}
        encoded_event = quote(event_type)
        encoded_prop = quote(prop["name"])
        update_result = _request(
            "PUT",
            f"taxonomy/event-property/{encoded_event}/{encoded_prop}",
            update_data,
            auth_header,
            dry_run,
            stats,
        )
        if update_result:
            stats["props_updated"] += 1
    elif result:
        stats["props_created"] += 1


def _get_catalog() -> list[dict[str, Any]]:
    """Load the event catalog from the bundled JSON file.

    The JSON is generated from taxonomy_catalog.py by generate_catalog_artifacts.py
    and shipped as package data so the SDK has no dependency on langley.
    """
    catalog_path = Path(resources.files("amplitude_ai") / "data" / "agent_event_catalog.json")  # type: ignore[arg-type]
    if not catalog_path.exists():
        logger.error(
            "agent_event_catalog.json not found at %s. "
            "Ensure the package was installed correctly.",
            catalog_path,
        )
        sys.exit(1)
    data = json.loads(catalog_path.read_text())
    events: list[dict[str, Any]] = data["events"]
    return events


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Register the [Agent] event schema in your Amplitude project's data catalog. "
            "Requires an Enterprise plan for Taxonomy API access."
        ),
    )
    parser.add_argument("--api-key", required=False, help="Amplitude project API key.")
    parser.add_argument("--secret-key", required=False, help="Amplitude project Secret key.")
    parser.add_argument("--dry-run", action="store_true", help="Preview without making API calls.")
    parser.add_argument("--eu", action="store_true", help="Use EU residency endpoint.")
    args = parser.parse_args()

    if args.eu:
        global BASE_URL
        BASE_URL = "https://analytics.eu.amplitude.com/api/2"

    if not args.dry_run and (not args.api_key or not args.secret_key):
        logger.error(
            "Provide --api-key and --secret-key, or use --dry-run to preview.\n"
            "Find your keys at: Amplitude > Settings > Projects > your project."
        )
        sys.exit(1)

    api_key = args.api_key or ""
    secret_key = args.secret_key or ""
    auth_header = _make_auth_header(api_key, secret_key)
    catalog = _get_catalog()

    stats: dict[str, int] = {
        "events_created": 0,
        "events_updated": 0,
        "props_created": 0,
        "props_updated": 0,
        "errors": 0,
    }

    _ensure_category(auth_header, args.dry_run, stats)

    total = len(catalog)
    for i, event in enumerate(catalog, 1):
        event_type = event["event_type"]
        props = event["properties"]
        logger.info("[%d/%d] %s (%d properties)", i, total, event_type, len(props))
        _register_event(event_type, event["description"], auth_header, args.dry_run, stats)
        for prop in props:
            _register_property(event_type, prop, auth_header, args.dry_run, stats)

    logger.info("=" * 60)
    logger.info("Registration Summary:")
    logger.info("  Events created:     %d", stats["events_created"])
    logger.info("  Events updated:     %d", stats["events_updated"])
    logger.info("  Properties created: %d", stats["props_created"])
    logger.info("  Properties updated: %d", stats["props_updated"])
    logger.info("  Errors:             %d", stats["errors"])
    logger.info("=" * 60)

    if args.dry_run:
        logger.info(
            "Dry run complete. No API calls were made. "
            "Events/properties show as 'created'; existing ones will be updated instead."
        )
    elif stats["errors"] > 0:
        logger.error("Completed with %d errors.", stats["errors"])
        sys.exit(1)
    else:
        logger.info("All events and properties registered successfully.")


if __name__ == "__main__":
    main()
