"""CLI for data registry management.

Commands for initializing, registering, and verifying versioned data artifacts.

Usage:
    # Initialize registry from existing files
    python -m embergraphrouter.versioning.cli init

    # Register a new dataset version
    python -m embergraphrouter.versioning.cli register dataset router_data data/router_data.csv

    # Verify integrity of all artifacts
    python -m embergraphrouter.versioning.cli verify

    # List all registered artifacts
    python -m embergraphrouter.versioning.cli list
"""

from __future__ import annotations

import argparse
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import yaml

from embergraphrouter.versioning.registry import (
    DataRegistry,
    compute_file_checksum,
    count_csv_rows,
    _MODULE_DIR,
    _REGISTRY_PATH,
)


def init_registry(force: bool = False) -> None:
    """Initialize registry from existing files.

    Scans data/ and configs/ directories to build initial registry.yaml.
    """
    if _REGISTRY_PATH.exists() and not force:
        print(f"Registry already exists at {_REGISTRY_PATH}")
        print("Use --force to overwrite.")
        return

    print("Initializing registry from existing files...")

    data_dir = _MODULE_DIR / "data"
    configs_dir = _MODULE_DIR / "configs"

    registry_data = {
        "schema_version": "1.0",
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "schemas": {
            "router_data_v2": {
                "required_columns": [
                    "task_id",
                    "task_description",
                    "query",
                    "ground_truth",
                    "metric",
                    "llm",
                    "effect",
                    "cost",
                    "latency",
                ],
                "optional_columns": [
                    "task_description_embedding",
                    "query_embedding",
                ],
                "embedding_dim": 256,
            },
            "unified_qa_v1": {
                "required_columns": [
                    "task_id",
                    "query",
                    "ground_truth",
                    "metric",
                ],
                "optional_columns": [],
            },
            "router_data_enriched_v1": {
                "required_columns": [
                    "task_id",
                    "task_description",
                    "query",
                    "ground_truth",
                    "metric",
                    "llm",
                    "effect",
                    "cost",
                    "latency",
                    "capability_primary",
                    "difficulty_band",
                ],
                "optional_columns": [
                    "task_description_embedding",
                    "query_embedding",
                ],
            },
        },
        "datasets": {},
        "configs": {},
    }

    # Register main datasets
    datasets_to_register = [
        ("router_data", "router_data.csv", "router_data_v2",
         "Main router training dataset with LLM responses and embeddings"),
        ("router_data_enriched", "router_data_enriched.csv", "router_data_enriched_v1",
         "Enriched router data with capability taxonomy and difficulty bands"),
        ("router_data_enriched_parquet", "router_data_enriched.parquet", "router_data_enriched_v1",
         "Parquet format of enriched router data for fast columnar reads"),
        ("unified_qa_data", "unified_qa_data.csv", "unified_qa_v1",
         "Unified QA dataset with tasks from MMLU, BBEH, LiveCodeBench, etc."),
    ]

    for name, filename, schema, description in datasets_to_register:
        filepath = data_dir / filename
        if filepath.exists():
            print(f"  Registering dataset: {name}")
            row_count = _count_rows(filepath)
            checksum = compute_file_checksum(filepath)
            mtime = datetime.fromtimestamp(filepath.stat().st_mtime, tz=timezone.utc)
            version = mtime.strftime("%Y.%m.%d")

            registry_data["datasets"][name] = {
                "description": description,
                "current_version": version,
                "schema": schema,
                "versions": {
                    version: {
                        "file": filename,
                        "checksum": checksum,
                        "row_count": row_count,
                        "created_at": mtime.isoformat(),
                        "source_config": "configs/config.yaml",
                    }
                },
            }
        else:
            print(f"  Skipping {name}: file not found ({filepath})")

    # Register main config
    config_path = configs_dir / "config.yaml"
    if config_path.exists():
        print("  Registering config: training")
        checksum = compute_file_checksum(config_path)
        mtime = datetime.fromtimestamp(config_path.stat().st_mtime, tz=timezone.utc)

        registry_data["configs"]["training"] = {
            "description": "Main training configuration for router GNN",
            "current_version": "v7",
            "versions": {
                "v7": {
                    "file": "configs/config.yaml",
                    "checksum": checksum,
                    "created_at": mtime.isoformat(),
                    "notes": "GAT architecture with quality-only training option",
                }
            },
        }

    # Register LLM descriptions
    llm_desc_path = configs_dir / "LLM_Descriptions.json"
    if llm_desc_path.exists():
        print("  Registering config: llm_descriptions")
        checksum = compute_file_checksum(llm_desc_path)
        mtime = datetime.fromtimestamp(llm_desc_path.stat().st_mtime, tz=timezone.utc)

        registry_data["configs"]["llm_descriptions"] = {
            "description": "LLM catalog with pricing, capabilities, and descriptions",
            "current_version": "v1",
            "versions": {
                "v1": {
                    "file": "configs/LLM_Descriptions.json",
                    "checksum": checksum,
                    "created_at": mtime.isoformat(),
                    "notes": "7-model catalog: Gemini, Claude, GPT variants",
                }
            },
        }

    # Write registry
    _REGISTRY_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(_REGISTRY_PATH, "w") as f:
        yaml.dump(registry_data, f, default_flow_style=False, sort_keys=False)

    print(f"\nRegistry created at {_REGISTRY_PATH}")
    print(f"  Datasets: {len(registry_data['datasets'])}")
    print(f"  Configs: {len(registry_data['configs'])}")
    print(f"  Schemas: {len(registry_data['schemas'])}")


def _count_rows(path: Path) -> int:
    """Count rows in CSV or Parquet file.

    Uses pandas for CSVs to handle edge cases like quoted newlines correctly.
    """
    if path.suffix == ".parquet":
        import pyarrow.parquet as pq
        return pq.read_metadata(path).num_rows
    elif path.suffix == ".csv":
        # Use pandas for accurate row count (handles quoted newlines, etc.)
        import pandas as pd
        # Read only index to minimize memory
        return len(pd.read_csv(path, usecols=[0]))
    else:
        return 0


def verify_registry() -> int:
    """Verify integrity of all registered artifacts.

    Returns:
        Exit code (0 if all valid, 1 if errors found).
    """
    try:
        registry = DataRegistry.load()
    except FileNotFoundError as e:
        print(f"Error: {e}")
        return 1

    print("Verifying registry integrity...")
    errors = registry.verify_all()

    if not errors:
        print("All artifacts verified successfully.")
        return 0

    print("\nIntegrity errors found:")
    for artifact, artifact_errors in errors.items():
        print(f"\n  {artifact}:")
        for err in artifact_errors:
            print(f"    - {err}")
    return 1


def list_artifacts() -> None:
    """List all registered artifacts."""
    try:
        registry = DataRegistry.load()
    except FileNotFoundError as e:
        print(f"Error: {e}")
        return

    print("Registered Datasets:")
    print("-" * 60)
    for name in registry.list_datasets():
        info = registry.get_dataset_info(name)
        current = info.get_version()
        print(f"  {name}")
        print(f"    Current: {info.current_version}")
        print(f"    Rows: {current.row_count:,}")
        print(f"    Schema: {info.schema_name}")
        print(f"    Versions: {', '.join(sorted(info.versions.keys()))}")
        print()

    print("\nRegistered Configs:")
    print("-" * 60)
    for name in registry.list_configs():
        info = registry.get_config_info(name)
        print(f"  {name}")
        print(f"    Current: {info.current_version}")
        print(f"    Versions: {', '.join(sorted(info.versions.keys()))}")
        print()


def register_dataset(
    name: str,
    filepath: str,
    schema: Optional[str] = None,
    version: Optional[str] = None,
    description: Optional[str] = None,
) -> None:
    """Register a new dataset or version."""
    path = Path(filepath)
    if not path.is_absolute():
        path = _MODULE_DIR / "data" / filepath

    if not path.exists():
        print(f"Error: File not found: {path}")
        return

    # Load existing registry
    try:
        with open(_REGISTRY_PATH) as f:
            registry_data = yaml.safe_load(f)
    except FileNotFoundError:
        print("Error: Registry not initialized. Run 'init' first.")
        return

    # Compute metadata
    row_count = _count_rows(path)
    checksum = compute_file_checksum(path)
    mtime = datetime.fromtimestamp(path.stat().st_mtime, tz=timezone.utc)
    version = version or mtime.strftime("%Y.%m.%d")
    relative_path = path.relative_to(_MODULE_DIR / "data")

    if name not in registry_data["datasets"]:
        # New dataset
        registry_data["datasets"][name] = {
            "description": description or f"Dataset: {name}",
            "current_version": version,
            "schema": schema or "",
            "versions": {},
        }
    else:
        # Update current version
        registry_data["datasets"][name]["current_version"] = version

    # Add version
    registry_data["datasets"][name]["versions"][version] = {
        "file": str(relative_path),
        "checksum": checksum,
        "row_count": row_count,
        "created_at": mtime.isoformat(),
    }

    registry_data["updated_at"] = datetime.now(timezone.utc).isoformat()

    with open(_REGISTRY_PATH, "w") as f:
        yaml.dump(registry_data, f, default_flow_style=False, sort_keys=False)

    print(f"Registered {name}@{version}")
    print(f"  Rows: {row_count:,}")
    print(f"  Checksum: {checksum[:40]}...")


def main() -> int:
    """CLI entry point."""
    parser = argparse.ArgumentParser(
        description="Data registry management for embergraphrouter"
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    # init command
    init_parser = subparsers.add_parser("init", help="Initialize registry from existing files")
    init_parser.add_argument("--force", "-f", action="store_true", help="Overwrite existing registry")

    # verify command
    subparsers.add_parser("verify", help="Verify integrity of all artifacts")

    # list command
    subparsers.add_parser("list", help="List all registered artifacts")

    # register command
    register_parser = subparsers.add_parser("register", help="Register a new dataset/config")
    register_parser.add_argument("type", choices=["dataset", "config"], help="Artifact type")
    register_parser.add_argument("name", help="Artifact name")
    register_parser.add_argument("path", help="Path to file")
    register_parser.add_argument("--schema", help="Schema name (datasets only)")
    register_parser.add_argument("--version", help="Version string (default: date-based)")
    register_parser.add_argument("--description", help="Human-readable description")

    args = parser.parse_args()

    if args.command == "init":
        init_registry(force=args.force)
        return 0
    elif args.command == "verify":
        return verify_registry()
    elif args.command == "list":
        list_artifacts()
        return 0
    elif args.command == "register":
        if args.type == "dataset":
            register_dataset(
                args.name,
                args.path,
                schema=args.schema,
                version=args.version,
                description=args.description,
            )
        else:
            print("Config registration not yet implemented")
        return 0

    return 1


if __name__ == "__main__":
    sys.exit(main())
