"""Multi-task graph router for LLM selection.

This module orchestrates the training pipeline for the EmberGraphRouter GNN,
handling data loading, preprocessing, splitting, and model training.
"""

from __future__ import annotations

import ast
import csv
import sys

csv.field_size_limit(sys.maxsize)
import json
import logging
import math
import random
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
import pandas as pd
import torch

from ..data_processing.utils import loadjson, loadpkl
from .graph_nn import FormData, GNNPredictor

if TYPE_CHECKING:
    import wandb as wandb_module

logger = logging.getLogger(__name__)

_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


class GraphRouterTrainer:
    """Orchestrates GNN training for LLM routing.

    This class handles:
    - Loading and parsing router data (queries, embeddings, metrics)
    - Stratified train/val/test splitting by task
    - Constructing PyG Data objects
    - Training the GNN with early stopping
    """

    def __init__(
        self,
        router_data_path: str,
        llm_path: str,
        llm_embedding_path: str,
        config: dict,
        wandb: wandb_module,
    ) -> None:
        self.config = config
        self.wandb = wandb
        self._router_data_path = router_data_path

        # Ensure model save directory exists
        model_path = Path(config["model_path"])
        if not model_path.is_absolute():
            model_path = Path.cwd() / model_path
        model_path.parent.mkdir(parents=True, exist_ok=True)
        self.config["model_path"] = str(model_path)

        # Load LLM descriptions
        self.llm_description = loadjson(llm_path)
        self.llm_names = list(self.llm_description.keys())
        self.num_llms = len(self.llm_names)
        self.config["llm_num"] = self.num_llms

        # Memory-efficient loading: only scalar columns first
        # Include 'query' for deduplication during train/val/test split
        logger.info("Loading scalar columns from router data...")
        scalar_cols = ["task_id", "effect", "cost", "latency", "llm", "query"]
        if str(router_data_path).endswith(".parquet"):
            self.data_df = pd.read_parquet(router_data_path, columns=scalar_cols)
        else:
            self.data_df = pd.read_csv(
                router_data_path,
                usecols=lambda c: c in scalar_cols,
            )
        total_rows = len(self.data_df)
        logger.info("Loaded %d rows (scalars only)", total_rows)

        # Group by (query, task_id) - each group should have exactly num_llms rows
        # This supports intentional oversampling where same query appears in different tasks
        self.data_df["_group_key"] = (
            self.data_df["query"].astype(str) + "|||" + self.data_df["task_id"].astype(str)
        )
        group_sizes = self.data_df.groupby("_group_key", sort=False).size()

        invalid_groups = group_sizes[group_sizes != self.num_llms]
        if len(invalid_groups) > 0:
            # Filter out invalid groups instead of failing
            valid_keys = group_sizes[group_sizes == self.num_llms].index
            original_rows = len(self.data_df)
            self.data_df = self.data_df[self.data_df["_group_key"].isin(valid_keys)]
            # Reset index so group_indices align with to_numpy() arrays
            self.data_df = self.data_df.reset_index(drop=True)
            logger.warning(
                "Filtered %d invalid (query, task_id) groups with wrong LLM count. "
                "Kept %d valid groups (%d rows -> %d rows, %.1f%% retained).",
                len(invalid_groups),
                len(valid_keys),
                original_rows,
                len(self.data_df),
                100 * len(self.data_df) / original_rows,
            )

        # Build group-to-row mapping for efficient lookup
        # group_indices[i] = list of row indices for the i-th unique (query, task_id) group
        self.group_indices: list[list[int]] = []
        self.query_task_labels: list[str] = []
        for group_key, group_df in self.data_df.groupby("_group_key", sort=False):
            self.group_indices.append(group_df.index.tolist())
            # task_id is the same for all rows in the group
            self.query_task_labels.append(group_df["task_id"].iloc[0])

        self.num_queries = len(self.group_indices)
        logger.info(
            "Grouped into %d unique (query, task_id) pairs (%d total rows)",
            self.num_queries,
            total_rows,
        )

        # Build task-to-query index
        self.task_to_query_indices: dict[str, list[int]] = {}
        for query_idx, task_id in enumerate(self.query_task_labels):
            self.task_to_query_indices.setdefault(task_id, []).append(query_idx)

        self.num_tasks = len(self.task_to_query_indices)
        self.config["num_task"] = self.num_tasks

        # Set random seed
        self._set_seed(config["seed"])

        # Load LLM embeddings
        self.llm_embeddings = loadpkl(llm_embedding_path)

        # Embedding column suffix for multi-embed parquet
        # e.g., "gemini_256" reads from "query_embedding_gemini_256"
        self._embedding_suffix = config.get("embedding_column_suffix", None)
        # Support concatenating multiple embeddings
        # e.g., ["gemini_768", "openai_large_256"] concatenates both
        self._embedding_suffixes = config.get("embedding_column_suffixes", None)

        # Prepare data (streams embeddings from CSV)
        self._prepare_data()
        self._split_data()

        # Initialize GNN
        self.form_data = FormData(_DEVICE)
        query_dim = self.query_embeddings.shape[1]
        llm_dim = self.llm_embeddings.shape[1]

        self.gnn = GNNPredictor(
            query_feature_dim=query_dim,
            llm_feature_dim=llm_dim,
            hidden_features_size=config["embedding_dim"],
            in_edges_size=config["edge_dim"],
            wandb=wandb,
            config=config,
            device=_DEVICE,
        )

        logger.info(
            "GNN initialized: %d queries, %d LLMs, %d tasks",
            self.num_queries,
            self.num_llms,
            self.num_tasks,
        )
        logger.info("Starting _train()...")
        self._train()

    def _set_seed(self, seed: int) -> None:
        """Set random seeds for reproducibility."""
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    def _parse_embedding(
        self, entry: str | list | np.ndarray, label: str, row: int
    ) -> list[float] | None:
        """Parse embedding into float list.

        Accepts strings (from CSV), native lists (from Parquet), or numpy
        arrays. Returns None for empty/invalid embeddings (caller should
        skip these rows).
        """
        # Native list or numpy array from Parquet — no parsing needed
        if isinstance(entry, np.ndarray):
            return entry.tolist() if entry.size > 0 else None
        if isinstance(entry, list):
            return entry if entry else None

        text = entry.strip()
        if not text or text == "[]":
            logger.debug("Empty %s embedding at row %d, skipping", label, row)
            return None

        # Try JSON/literal_eval first
        for parser in (json.loads, ast.literal_eval):
            try:
                parsed = parser(text)
                # Handle nested lists
                if isinstance(parsed, list) and parsed and isinstance(parsed[0], list):
                    return parsed[0]
                if isinstance(parsed, list):
                    return parsed
            except (json.JSONDecodeError, ValueError, SyntaxError):
                continue

        # Fallback: parse as space-separated floats
        vector = np.fromstring(text.replace("[", " ").replace("]", " "), sep=" ")
        if vector.size == 0:
            raise ValueError(f"Failed to parse {label} embedding at row {row}.")
        return vector.tolist()

    def _prepare_data(self) -> None:
        """Parse embeddings and metrics from dataframe.

        Memory optimization: Stream embeddings from CSV, only parsing the first
        row of each (query, task_id) group instead of loading the full 2.4GB file.
        """
        # Extract scalar metrics (already loaded efficiently)
        self.perf_list = self.data_df["effect"].to_numpy(dtype=np.float32)
        self.cost_list = self.data_df["cost"].to_numpy(dtype=np.float32)

        if "latency" in self.data_df.columns:
            self.latency_list = self.data_df["latency"].to_numpy(dtype=np.float32)
        else:
            self.latency_list = np.zeros(len(self.data_df), dtype=np.float32)

        # Get the first row index from each (query, task_id) group for embedding parsing
        # All rows in a group share the same query/task embeddings
        unique_row_indices = set(group[0] for group in self.group_indices)
        logger.info("Streaming embeddings for %d unique queries...", len(unique_row_indices))

        # Create a mapping from row index to group index for proper ordering
        row_to_group_idx = {
            group[0]: group_idx for group_idx, group in enumerate(self.group_indices)
        }

        query_embeddings_dict: dict[int, list[float] | None] = {}
        task_embeddings_dict: dict[int, list[float] | None] = {}

        # Handle both parquet and CSV formats for embedding streaming
        if str(self._router_data_path).endswith(".parquet"):
            # Parquet: load embedding columns directly
            # Support multi-embed parquet via embedding_column_suffix config
            # Also supports concatenating multiple embeddings via embedding_column_suffixes
            if self._embedding_suffixes:
                # Concatenate multiple embedding columns
                query_cols = [f"query_embedding_{s}" for s in self._embedding_suffixes]
                task_cols = [f"task_description_embedding_{s}" for s in self._embedding_suffixes]
                all_cols = query_cols + task_cols
                logger.info("Concatenating embeddings from: %s", self._embedding_suffixes)
                embed_df = pd.read_parquet(self._router_data_path, columns=all_cols)
                for row_idx in unique_row_indices:
                    group_idx = row_to_group_idx[row_idx]
                    # Concatenate query embeddings
                    query_parts = []
                    for col in query_cols:
                        emb = embed_df.iloc[row_idx][col]
                        parsed = self._parse_embedding(
                            emb if pd.notna(emb) else "", "query", row_idx
                        )
                        if parsed is not None:
                            query_parts.append(parsed)
                    # Concatenate task embeddings
                    task_parts = []
                    for col in task_cols:
                        emb = embed_df.iloc[row_idx][col]
                        parsed = self._parse_embedding(
                            emb if pd.notna(emb) else "", "task", row_idx
                        )
                        if parsed is not None:
                            task_parts.append(parsed)
                    # Store concatenated embeddings
                    if query_parts:
                        query_embeddings_dict[group_idx] = np.concatenate(query_parts)
                    else:
                        query_embeddings_dict[group_idx] = None
                    if task_parts:
                        task_embeddings_dict[group_idx] = np.concatenate(task_parts)
                    else:
                        task_embeddings_dict[group_idx] = None
                    if len(query_embeddings_dict) % 5000 == 0:
                        logger.info("  Parsed %d embeddings...", len(query_embeddings_dict))
            elif self._embedding_suffix:
                query_col = f"query_embedding_{self._embedding_suffix}"
                task_col = f"task_description_embedding_{self._embedding_suffix}"
                logger.info("Using embedding columns: %s, %s", query_col, task_col)
                embed_df = pd.read_parquet(
                    self._router_data_path,
                    columns=[query_col, task_col],
                )
                for row_idx in unique_row_indices:
                    group_idx = row_to_group_idx[row_idx]
                    query_emb = embed_df.iloc[row_idx][query_col]
                    task_emb = embed_df.iloc[row_idx][task_col]
                    # Pass directly — _parse_embedding handles lists,
                    # arrays, strings, and None values.
                    query_embeddings_dict[group_idx] = self._parse_embedding(
                        query_emb if query_emb is not None else "",
                        "query", row_idx,
                    )
                    task_embeddings_dict[group_idx] = self._parse_embedding(
                        task_emb if task_emb is not None else "",
                        "task", row_idx,
                    )
                    if len(query_embeddings_dict) % 5000 == 0:
                        logger.info("  Parsed %d embeddings...", len(query_embeddings_dict))
            else:
                query_col = "query_embedding"
                task_col = "task_description_embedding"
                embed_df = pd.read_parquet(
                    self._router_data_path,
                    columns=[query_col, task_col],
                )
                for row_idx in unique_row_indices:
                    group_idx = row_to_group_idx[row_idx]
                    query_emb = embed_df.iloc[row_idx][query_col]
                    task_emb = embed_df.iloc[row_idx][task_col]
                    query_embeddings_dict[group_idx] = self._parse_embedding(
                        query_emb if pd.notna(query_emb) else "", "query", row_idx
                    )
                    task_embeddings_dict[group_idx] = self._parse_embedding(
                        task_emb if pd.notna(task_emb) else "", "task", row_idx
                    )
                    if len(query_embeddings_dict) % 5000 == 0:
                        logger.info("  Parsed %d embeddings...", len(query_embeddings_dict))
        else:
            # CSV: stream with DictReader
            with open(self._router_data_path, "r", encoding="utf-8") as f:
                reader = csv.DictReader(f)
                for row_idx, row in enumerate(reader):
                    if row_idx in unique_row_indices:
                        group_idx = row_to_group_idx[row_idx]
                        query_embeddings_dict[group_idx] = self._parse_embedding(
                            row["query_embedding"], "query", row_idx
                        )
                        task_embeddings_dict[group_idx] = self._parse_embedding(
                            row["task_description_embedding"], "task", row_idx
                        )
                        if len(query_embeddings_dict) % 5000 == 0:
                            logger.info("  Parsed %d embeddings...", len(query_embeddings_dict))

        # Filter out groups with empty embeddings
        valid_group_indices = [
            i for i in range(len(self.group_indices))
            if query_embeddings_dict.get(i) is not None
            and task_embeddings_dict.get(i) is not None
        ]
        skipped = len(self.group_indices) - len(valid_group_indices)
        if skipped > 0:
            logger.warning(
                "Skipped %d groups with empty embeddings (%.1f%% of data)",
                skipped, 100 * skipped / len(self.group_indices),
            )
            # Rebuild group_indices and related mappings
            old_to_new = {old: new for new, old in enumerate(valid_group_indices)}
            self.group_indices = [self.group_indices[i] for i in valid_group_indices]
            self.query_task_labels = [self.query_task_labels[i] for i in valid_group_indices]
            self.num_queries = len(self.group_indices)
            # Rebuild task_to_query_indices
            self.task_to_query_indices = {}
            for query_idx, task_id in enumerate(self.query_task_labels):
                self.task_to_query_indices.setdefault(task_id, []).append(query_idx)

        # Convert to numpy arrays in proper order (only valid groups)
        query_embeddings = [
            query_embeddings_dict[i] for i in valid_group_indices
        ]
        task_embeddings = [
            task_embeddings_dict[i] for i in valid_group_indices
        ]

        self.query_embeddings = np.array(query_embeddings, dtype=np.float32)
        self.task_embeddings = np.array(task_embeddings, dtype=np.float32)

        logger.info("Query embeddings shape: %s", self.query_embeddings.shape)

    def _reorder_row_data(self) -> None:
        """Reorder row-level data arrays to match group ordering.

        The data was loaded in CSV row order, but we need it ordered by
        (query, task_id) groups for the GNN edge indexing to work correctly.
        After reordering, edge i maps to reordered row i.
        """
        # Flatten group_indices to get all row indices in group order
        reorder_indices = []
        for group in self.group_indices:
            reorder_indices.extend(group)

        reorder_indices = np.array(reorder_indices)

        # Reorder all row-level arrays
        self.perf_list = self.perf_list[reorder_indices]
        self.cost_list = self.cost_list[reorder_indices]
        self.latency_list = self.latency_list[reorder_indices]

        logger.info("Reordered row data to match group ordering (%d rows)", len(reorder_indices))

    def _compute_split(self, count: int) -> tuple[int, int, int]:
        """Compute train/val/test split ensuring all sets are non-empty."""
        split_ratio = self.config["split_ratio"]

        exact = [count * ratio for ratio in split_ratio]
        floors = [int(math.floor(v)) for v in exact]
        remainder = count - sum(floors)

        # Distribute remainder to largest fractional parts
        fractional = sorted(
            ((exact[i] - floors[i], i) for i in range(len(split_ratio))),
            reverse=True,
        )
        for idx in range(remainder):
            _, target = fractional[idx % len(split_ratio)]
            floors[target] += 1

        train_c, val_c = floors[0], floors[1]
        test_c = count - train_c - val_c

        # Ensure all splits are non-empty for sufficient data
        if test_c == 0 and count > 0:
            if train_c > 1:
                train_c -= 1
                test_c += 1
            elif val_c > 0:
                val_c -= 1
                test_c += 1

        if val_c == 0 and count >= 3:
            if train_c > 1:
                train_c -= 1
                val_c += 1
            elif test_c > 1:
                test_c -= 1
                val_c += 1

        if train_c == 0 and count >= 3:
            if test_c > 1:
                test_c -= 1
                train_c += 1
            elif val_c > 1:
                val_c -= 1
                train_c += 1

        if train_c < 0 or val_c < 0 or test_c < 0:
            raise ValueError("Invalid split calculation.")
        if train_c + val_c + test_c != count:
            raise ValueError("Split counts do not sum to total.")

        return train_c, val_c, test_c

    def _split_data(self) -> None:
        """Create stratified train/val/test masks with query deduplication.

        Ensures that duplicate queries (same text appearing multiple times due
        to upsampling) are assigned to the same split, preventing data leakage
        between train/val/test sets.

        With (query, task_id) grouping, edges map to group indices rather than
        raw row indices. Each group has num_llms edges.
        """
        # Reorder data arrays to match group ordering
        # This flattens group_indices to get all row indices in group order
        self._reorder_row_data()

        # Combine edge features (now in group order)
        self.combined_edge = np.column_stack([
            self.cost_list,
            self.perf_list,
            self.latency_list,
        ])

        self.labels = np.zeros((len(self.perf_list), 1), dtype=float)

        # Build edge indices (now contiguous by group)
        self.source_nodes = [q for q in range(self.num_queries) for _ in range(self.num_llms)]
        self.dest_nodes = list(range(self.num_llms)) * self.num_queries

        total_edges = len(self.source_nodes)

        # Initialize masks
        self.train_mask = torch.zeros(total_edges, dtype=torch.bool)
        self.valid_mask = torch.zeros(total_edges, dtype=torch.bool)
        self.test_mask = torch.zeros(total_edges, dtype=torch.bool)

        # Extract query texts for deduplication using first row of each group
        query_texts = [
            self.data_df.loc[self.group_indices[i][0], "query"]
            for i in range(self.num_queries)
        ]

        # Stratified split by task with deduplication
        rng = random.Random(self.config["seed"])
        total_unique = 0
        total_duplicated = 0

        for task_id, query_indices in self.task_to_query_indices.items():
            # Group query indices by unique query text
            # This ensures all duplicates of the same query go to the same split
            unique_query_to_indices: dict[str, list[int]] = {}
            for query_idx in query_indices:
                query_text = query_texts[query_idx]
                unique_query_to_indices.setdefault(query_text, []).append(query_idx)

            # Shuffle unique queries (not individual indices)
            unique_queries = list(unique_query_to_indices.keys())
            rng.shuffle(unique_queries)

            # Compute split based on unique query count
            train_c, val_c, test_c = self._compute_split(len(unique_queries))

            train_unique = unique_queries[:train_c]
            val_unique = unique_queries[train_c:train_c + val_c]
            test_unique = unique_queries[train_c + val_c:]

            # Map unique queries back to all their indices
            # Edge indices are now contiguous: group i spans [i*num_llms, (i+1)*num_llms)
            for mask, unique_list in [
                (self.train_mask, train_unique),
                (self.valid_mask, val_unique),
                (self.test_mask, test_unique),
            ]:
                for unique_query in unique_list:
                    for query_idx in unique_query_to_indices[unique_query]:
                        start = query_idx * self.num_llms
                        end = start + self.num_llms
                        mask[start:end] = True

            # Track deduplication stats
            task_unique = len(unique_queries)
            task_total = len(query_indices)
            total_unique += task_unique
            total_duplicated += task_total - task_unique

        logger.info(
            "Split complete: %d unique queries, %d duplicates removed from split calculation",
            total_unique, total_duplicated
        )

    def _train(self) -> None:
        """Build PyG data objects and train the GNN."""
        logger.info("Building train_data...")
        # Training data
        train_data = self.form_data.formulation(
            task_features=self.task_embeddings,
            query_features=self.query_embeddings,
            llm_features=self.llm_embeddings,
            source_nodes=self.source_nodes,
            dest_nodes=self.dest_nodes,
            edge_features=self.perf_list,
            labels=self.labels,
            edge_mask=self.train_mask,
            combined_edge=self.combined_edge,
            train_mask=self.train_mask,
            valid_mask=self.valid_mask,
            test_mask=self.test_mask,
            query_task_labels=self.query_task_labels,
        )

        logger.info("train_data built. Building valid_data...")
        # Validation data
        valid_data = self.form_data.formulation(
            task_features=self.task_embeddings,
            query_features=self.query_embeddings,
            llm_features=self.llm_embeddings,
            source_nodes=self.source_nodes,
            dest_nodes=self.dest_nodes,
            edge_features=self.perf_list,
            labels=self.labels,
            edge_mask=self.valid_mask,
            combined_edge=self.combined_edge,
            train_mask=self.train_mask,
            valid_mask=self.valid_mask,
            test_mask=self.test_mask,
            query_task_labels=self.query_task_labels,
        )

        logger.info("valid_data built. Building test_data...")
        # Test data
        test_data = self.form_data.formulation(
            task_features=self.task_embeddings,
            query_features=self.query_embeddings,
            llm_features=self.llm_embeddings,
            source_nodes=self.source_nodes,
            dest_nodes=self.dest_nodes,
            edge_features=self.perf_list,
            labels=self.labels,
            edge_mask=self.test_mask,
            combined_edge=self.combined_edge,
            train_mask=self.train_mask,
            valid_mask=self.valid_mask,
            test_mask=self.test_mask,
            query_task_labels=self.query_task_labels,
        )

        self._test_data = test_data

        logger.info("All data built. Starting gnn.train_validate()...")
        # Train
        self.gnn.train_validate(
            data=train_data,
            data_validate=valid_data,
            data_for_test=test_data,
        )

        # Post-hoc temperature calibration for better probability estimates
        logger.info("Running post-hoc temperature calibration...")
        self.gnn.calibrate_temperature(valid_data)

    def route_with(
        self,
        perf: float,
        cost: float,
        latency: float,
    ) -> tuple[dict[str, float], float]:
        """Evaluate router with custom preference weights.

        Args:
            perf: Weight for performance (quality).
            cost: Weight for cost.
            latency: Weight for latency.

        Returns:
            Tuple of (metrics dict, loss).
        """
        previous = (self.gnn.infer_wp, self.gnn.infer_wc, self.gnn.infer_wl)
        self.gnn.infer_wp = float(perf)
        self.gnn.infer_wc = float(cost)
        self.gnn.infer_wl = float(latency)

        try:
            return self.gnn.test(
                data=self._test_data,
                model_path=self.config["model_path"],
            )
        finally:
            self.gnn.infer_wp, self.gnn.infer_wc, self.gnn.infer_wl = previous

    def test(self) -> tuple[dict[str, float], float]:
        """Evaluate router with default inference weights."""
        return self.gnn.test(
            data=self._test_data,
            model_path=self.config["model_path"],
        )


# Backwards compatibility alias
graph_router_prediction = GraphRouterTrainer


if __name__ == "__main__":
    import argparse
    import sys

    import wandb
    import yaml

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    parser = argparse.ArgumentParser(description="Train GNN router")
    parser.add_argument(
        "--config", default="config.yaml", help="Config filename in configs/"
    )
    args = parser.parse_args()

    base_dir = Path(__file__).resolve().parents[1]

    def _resolve(path_str: str) -> str:
        candidate = Path(path_str)
        if candidate.is_absolute():
            return str(candidate)
        return str((base_dir / candidate).resolve())

    config_path = base_dir / "configs" / args.config
    with config_path.open("r", encoding="utf-8") as f:
        config = yaml.safe_load(f)

    wandb_key = config.get("wandb_key")
    if wandb_key:
        wandb.login(key=wandb_key)
        wandb.init(project="graph_router")
    else:
        wandb.init(project="graph_router", mode="disabled")

    GraphRouterTrainer(
        router_data_path=_resolve(config["saved_router_data_path"]),
        llm_path=_resolve(config["llm_description_path"]),
        llm_embedding_path=_resolve(config["llm_embedding_path"]),
        config=config,
        wandb=wandb,
    )
