#!/usr/bin/env python3
"""Upload enriched router data to GCS with proper organization.

Data Organization (Dean-style):
    gs://eb1-experiment-backups/
    └── datasets/
        └── graphrouter/
            └── v1/                          # Schema version
                └── 2025-12-14/              # Snapshot date
                    ├── manifest.json        # Metadata, schema, stats
                    ├── router_data.parquet  # Main training data
                    ├── router_data.csv      # Human-readable copy
                    ├── cache/
                    │   └── response_cache.json
                    └── configs/
                        └── LLM_Descriptions.json

Naming Conventions:
    - Versions are MAJOR only (v1, v2) - bump on schema changes
    - Dates are ISO format (2025-12-14)
    - Underscores for multi-word file names
    - No spaces, no special characters

Usage:
    python -m embergraphrouter.scripts.upload_to_gcs
    python -m embergraphrouter.scripts.upload_to_gcs --version v2  # New schema
    python -m embergraphrouter.scripts.upload_to_gcs --dry-run     # Preview only
"""

import argparse
import json
import logging
import subprocess
from datetime import datetime
from pathlib import Path

import pandas as pd

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

_MODULE_DIR = Path(__file__).resolve().parent.parent

# GCS Configuration
GCS_BUCKET = "gs://eb1-experiment-backups"
GCS_PREFIX = "datasets/graphrouter"


def generate_manifest(
    data_path: Path,
    cache_path: Path,
    llm_desc_path: Path,
    version: str,
    date: str,
) -> dict:
    """Generate manifest.json with dataset metadata."""

    df = pd.read_parquet(data_path)

    # Schema documentation
    schema = {
        col: {
            "dtype": str(df[col].dtype),
            "description": get_column_description(col),
            "sample": str(df[col].iloc[0])[:100] if len(df) > 0 else None,
        }
        for col in df.columns
    }

    # Statistics
    stats = {
        "total_rows": len(df),
        "total_queries": df["query"].nunique(),
        "total_models": df["llm"].nunique(),
        "models": df["llm"].unique().tolist(),
        "tasks": df["task_id"].unique().tolist(),
        "response_types": df["response_type"].value_counts().to_dict(),
        "effect_by_model": df.groupby("llm")["effect"].mean().to_dict(),
        "effect_by_tier": df.groupby("model_tier")["effect"].mean().to_dict(),
        "effect_by_capability": df.groupby("capability_primary")["effect"].mean().to_dict(),
    }

    # File info
    files = {
        "router_data.parquet": {
            "rows": len(df),
            "columns": len(df.columns),
            "size_mb": data_path.stat().st_size / (1024 * 1024),
        },
    }

    if cache_path.exists():
        cache = json.loads(cache_path.read_text())
        files["cache/response_cache.json"] = {
            "entries": len(cache),
            "size_mb": cache_path.stat().st_size / (1024 * 1024),
        }

    return {
        "version": version,
        "date": date,
        "created_at": datetime.utcnow().isoformat() + "Z",
        "description": "GraphRouter training dataset with enriched labels",
        "schema": schema,
        "statistics": stats,
        "files": files,
        "provenance": {
            "source": "embergraphrouter.data_processing",
            "scripts": [
                "multidata_unify.py",
                "construct_router_data.py",
                "enrich_router_data.py",
            ],
        },
    }


def get_column_description(col: str) -> str:
    """Return human-readable description for each column."""
    descriptions = {
        "task_id": "Source dataset identifier (e.g., mmlu, arc, GSM8K)",
        "query": "Input prompt text",
        "llm": "Model name used for this response",
        "capability_primary": "Primary capability type: knowledge, reasoning, generation, code",
        "domain": "Subject domain: science, law, math, language, etc.",
        "capability_tags": "Comma-separated fine-grained capability tags",
        "model_tier": "Cost tier: budget, standard, premium",
        "model_provider": "API provider: openai, anthropic, google",
        "input_price_per_m": "Input token price per million tokens (USD)",
        "output_price_per_m": "Output token price per million tokens (USD)",
        "effect": "Binary correctness: 1.0 if correct, 0.0 if wrong",
        "response_type": "Response classification: success, refusal, empty",
        "latency_sec": "Actual response latency in seconds",
        "cost": "Normalized cost (0-1 within task)",
        "latency": "Normalized latency (0-1 within task)",
        "query_difficulty": "Avg effect across all models (0=hard, 1=easy)",
        "model_advantage": "This model's effect minus query_difficulty",
        "task_difficulty": "Avg effect across all queries in task",
        "model_task_affinity": "Model's performance on this task vs overall",
        "ground_truth": "Expected correct answer",
        "metric": "Evaluation metric: em_mc, f1_score, GSM8K, etc.",
        "task_description": "Natural language description of the task",
        "query_embedding": "256-dim embedding of query text",
        "task_description_embedding": "256-dim embedding of task description",
    }
    return descriptions.get(col, "No description available")


def gsutil_cp(src: Path, dst: str, dry_run: bool = False) -> bool:
    """Copy file to GCS using gsutil."""
    cmd = ["gsutil", "-q", "cp", str(src), dst]
    if dry_run:
        logger.info(f"[DRY RUN] {' '.join(cmd)}")
        return True

    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        logger.error(f"Failed: {result.stderr}")
        return False
    return True


def upload_dataset(
    version: str = "v1",
    dry_run: bool = False,
) -> None:
    """Upload enriched router data to GCS."""

    date = datetime.now().strftime("%Y-%m-%d")
    gcs_base = f"{GCS_BUCKET}/{GCS_PREFIX}/{version}/{date}"

    logger.info(f"Uploading to: {gcs_base}")
    logger.info(f"Version: {version}, Date: {date}")

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

    # Files to upload
    files = [
        (data_dir / "router_data_enriched.parquet", f"{gcs_base}/router_data.parquet"),
        (data_dir / "router_data_enriched.csv", f"{gcs_base}/router_data.csv"),
        (data_dir / ".router_cache" / "response_cache.json", f"{gcs_base}/cache/response_cache.json"),
        (configs_dir / "LLM_Descriptions.json", f"{gcs_base}/configs/LLM_Descriptions.json"),
    ]

    # Generate and upload manifest first
    manifest_path = data_dir / "manifest.json"
    manifest = generate_manifest(
        data_path=data_dir / "router_data_enriched.parquet",
        cache_path=data_dir / ".router_cache" / "response_cache.json",
        llm_desc_path=configs_dir / "LLM_Descriptions.json",
        version=version,
        date=date,
    )

    manifest_path.write_text(json.dumps(manifest, indent=2))
    files.insert(0, (manifest_path, f"{gcs_base}/manifest.json"))

    # Upload all files
    success_count = 0
    for src, dst in files:
        if not src.exists():
            logger.warning(f"Skipping missing file: {src}")
            continue

        size_mb = src.stat().st_size / (1024 * 1024)
        logger.info(f"Uploading {src.name} ({size_mb:.1f} MB)...")

        if gsutil_cp(src, dst, dry_run):
            success_count += 1
        else:
            logger.error(f"Failed to upload: {src}")

    logger.info(f"\nUploaded {success_count}/{len(files)} files to {gcs_base}")

    # Print the GCS paths for easy access
    print("\n" + "=" * 60)
    print("GCS PATHS")
    print("=" * 60)
    print(f"\nBase path: {gcs_base}")
    print(f"\nManifest:  {gcs_base}/manifest.json")
    print(f"Data:      {gcs_base}/router_data.parquet")
    print(f"\nTo download:")
    print(f"  gsutil cp {gcs_base}/manifest.json .")
    print(f"  gsutil cp {gcs_base}/router_data.parquet .")

    # Clean up local manifest
    if manifest_path.exists():
        manifest_path.unlink()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Upload router data to GCS")
    parser.add_argument(
        "--version",
        type=str,
        default="v1",
        help="Schema version (e.g., v1, v2). Bump on schema changes.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Preview upload without actually uploading",
    )
    args = parser.parse_args()

    upload_dataset(version=args.version, dry_run=args.dry_run)
