"""Train the EmberGraphRouter GNN model with post-training evaluation."""

import argparse
import logging
import sys
from pathlib import Path

# Ensure parent directory is on path for relative imports
_MODULE_DIR = Path(__file__).resolve().parent
if str(_MODULE_DIR.parent) not in sys.path:
    sys.path.insert(0, str(_MODULE_DIR.parent))

import wandb
import yaml

from embergraphrouter.versioning import DataRegistry
from embergraphrouter.evaluation import run_post_training_eval, PostTrainingEvalConfig

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger(__name__)

from embergraphrouter.model.multi_task_graph_router import graph_router_prediction

_EMBER_CONFIG_PATH = Path.home() / ".ember" / "config.yaml"


def _load_wandb_key(local_config: dict) -> str | None:
    """Load W&B API key from local config or ~/.ember/config.yaml."""
    # First check local config
    if local_config.get("wandb_key"):
        return local_config["wandb_key"]

    # Fall back to ~/.ember/config.yaml
    if _EMBER_CONFIG_PATH.exists():
        with open(_EMBER_CONFIG_PATH, "r", encoding="utf-8") as f:
            ember_config = yaml.safe_load(f) or {}
        wandb_section = ember_config.get("wandb", {})
        if isinstance(wandb_section, dict):
            return wandb_section.get("api_key")

    return None


def main():
    parser = argparse.ArgumentParser(description="Train EmberGraphRouter GNN")
    parser.add_argument(
        "--config_file",
        type=str,
        default=str(_MODULE_DIR / "configs" / "config.yaml"),
        help="Path to config YAML file",
    )
    parser.add_argument(
        "--offline",
        action="store_true",
        help="Run W&B in offline mode",
    )
    parser.add_argument(
        "--use-versioning",
        action="store_true",
        help="Use versioned data registry (auto-selects Parquet, validates checksums)",
    )
    parser.add_argument(
        "--dataset",
        type=str,
        default="router_data",
        help="Dataset name from registry (default: router_data)",
    )
    # Sweep parameters (override config values)
    parser.add_argument("--gnn_architecture", type=str, help="GNN architecture (gcn/gat)")
    parser.add_argument("--embedding_dim", type=int, help="Embedding dimension")
    parser.add_argument("--num_gnn_layers", type=int, help="Number of GNN layers")
    parser.add_argument("--difficulty_loss_weight", type=float, help="Difficulty loss weight")

    # Post-training evaluation
    parser.add_argument(
        "--eval",
        action="store_true",
        default=True,
        help="Run post-training evaluation (default: True)",
    )
    parser.add_argument(
        "--no-eval",
        action="store_true",
        help="Skip post-training evaluation",
    )
    parser.add_argument(
        "--eval-advanced",
        action="store_true",
        default=True,
        help="Run advanced evaluation (error, cost, significance analysis)",
    )

    args = parser.parse_args()

    with open(args.config_file, "r", encoding="utf-8") as file:
        config = yaml.safe_load(file)

    # Resolve data paths - either via versioning or direct paths
    if args.use_versioning:
        logger.info("Using versioned data registry (Parquet preferred)")
        registry = DataRegistry.load()

        # Get dataset path from registry (auto-selects Parquet)
        dataset_name = args.dataset
        parquet_map = {
            "router_data": "router_data_parquet",
            "router_data_enriched": "router_data_enriched_parquet",
        }
        if dataset_name in parquet_map:
            try:
                registry.get_dataset_info(parquet_map[dataset_name])
                dataset_name = parquet_map[dataset_name]
                logger.info(f"Auto-selected Parquet: {args.dataset} -> {dataset_name}")
            except KeyError:
                pass

        dataset_info = registry.get_dataset_info(dataset_name)
        version = dataset_info.get_version()
        config["saved_router_data_path"] = str(version.absolute_path)
        logger.info(f"Dataset: {dataset_name}@{version.version} ({version.row_count:,} rows)")

        # Resolve other paths normally
        for key in ("unified_qa_data_path", "llm_description_path", "llm_embedding_path", "model_path"):
            if key in config and not Path(config[key]).is_absolute():
                config[key] = str(_MODULE_DIR / config[key])
    else:
        # Legacy: resolve relative paths to absolute
        for key in (
            "unified_qa_data_path",
            "llm_description_path",
            "saved_router_data_path",
            "llm_embedding_path",
            "model_path",
        ):
            if key in config and not Path(config[key]).is_absolute():
                config[key] = str(_MODULE_DIR / config[key])

    # Initialize W&B
    wandb_key = _load_wandb_key(config)
    if wandb_key:
        wandb.login(key=wandb_key)
        mode = "offline" if args.offline else "online"
        wandb.init(project="graph_router", mode=mode, config=config)
    else:
        logger.warning("No W&B key found; running in disabled mode")
        wandb.init(project="graph_router", mode="disabled", config=config)

    # Override config with command-line arguments (for WandB sweeps)
    for param in ["gnn_architecture", "embedding_dim", "num_gnn_layers", "difficulty_loss_weight"]:
        value = getattr(args, param, None)
        if value is not None:
            logger.info(f"CLI override: {param}={config.get(param, 'N/A')} -> {value}")
            config[param] = value

    # Also check WandB config for any additional overrides
    if wandb.run is not None and wandb.config:
        for key, value in wandb.config.items():
            if key in config and key not in ["gnn_architecture", "embedding_dim", "num_gnn_layers", "difficulty_loss_weight"]:
                logger.info(f"WandB config override: {key}={config[key]} -> {value}")
                config[key] = value

    logger.info("Starting GraphRouterTrainer...")
    graph_router_prediction(
        router_data_path=config["saved_router_data_path"],
        llm_path=config["llm_description_path"],
        llm_embedding_path=config["llm_embedding_path"],
        config=config,
        wandb=wandb,
    )

    # Post-training evaluation
    if args.eval and not args.no_eval:
        logger.info("Running post-training evaluation...")
        eval_config = PostTrainingEvalConfig(
            enabled=True,
            run_advanced=args.eval_advanced,
            print_scorecard=True,
            save_results=True,
        )
        eval_results = run_post_training_eval(
            model_path=config["model_path"],
            data_path=config["saved_router_data_path"],
            config=config,
            eval_config=eval_config,
            llm_embedding_path=config.get("llm_embedding_path"),
        )
        if eval_results:
            # Log to W&B
            wandb.log({
                "final/frac_oracle": eval_results.frac_oracle,
                "final/top1_accuracy": eval_results.top1_accuracy,
                "final/ece": eval_results.ece,
                "final/difficulty_auc": eval_results.difficulty_auc,
                "final/mean_logit_corr": eval_results.mean_logit_corr,
                "final/spearman_corr": eval_results.spearman_corr,
            })
            logger.info("Post-training evaluation complete. Results: %s", eval_results.results_path)
        else:
            logger.warning("Post-training evaluation returned no results")
    else:
        logger.info("Skipping post-training evaluation (--no-eval)")

    wandb.finish()


if __name__ == "__main__":
    main()
