from __future__ import annotations

import logging
import re
from typing import Any

from expops.core.runtime.step_system import (
    ProcessDefinition,
    get_current_context,
    get_current_process_context,
    get_process_registry,
    log_metric,
    process,
)

from .registry import get_adapter, get_utility

COMPONENT_PREFIX = "sklearn."
COMPONENT_RUNNER_NAME = "run_sklearn_component"
_RESERVED_PARAM_KEYS = {
    "operation",
}
logger = logging.getLogger(__name__)
 
 
def _base_process_name(name: str) -> str:
    base = str(name)
    seed_match = re.match(r"^(.*)_seed-?\d+$", base)
    if seed_match and seed_match.group(1):
        base = seed_match.group(1)
    part_match = re.match(r"^(.*)__p\d+$", base)
    if part_match and part_match.group(1):
        base = part_match.group(1)
    return base
 
 
def _lookup_process_cfg(context: Any, process_name: str) -> dict[str, Any]:
    if not context or not isinstance(getattr(context, "global_config", None), dict):
        return {}
    pipeline_cfg = context.global_config.get("pipeline") or {}
    processes = pipeline_cfg.get("processes") if isinstance(pipeline_cfg, dict) else None
    if not isinstance(processes, list):
        return {}
    base_name = _base_process_name(process_name)
    names_to_match = {process_name, base_name}
    if process_name.startswith("//") and "process[@name=" in process_name:
        try:
            from expops.core.graph.pipeline_tree import _extract_base_name_from_canonical_xpath

            xpath_base = _extract_base_name_from_canonical_xpath(process_name)
            if xpath_base:
                names_to_match.add(xpath_base)
        except Exception:
            pass
    for proc in processes:
        proc_name = proc.get("name") if isinstance(proc, dict) else getattr(proc, "name", None)
        if proc_name is not None and proc_name in names_to_match:
            if isinstance(proc, dict):
                return proc
            try:
                keys = (
                    "name",
                    "component",
                    "parameters",
                    "input_transform",
                    "output_mapping",
                    "metrics",
                    "environment",
                    "iterative_logging",
                    "inputs",
                    "outputs",
                )
                return {k: getattr(proc, k, None) for k in keys}
            except Exception:
                return {}
    return {}
 
 
def _alias_map(cfg: dict[str, Any], key: str) -> dict[str, str]:
    raw = cfg.get(key)
    if not isinstance(raw, dict):
        return {}
    out: dict[str, str] = {}
    for canonical, alias in raw.items():
        if canonical is None or alias is None:
            continue
        out[str(canonical)] = str(alias)
    return out
 
 
def _resolve_input(
    available: dict[str, Any],
    aliases: dict[str, str],
    canonical_key: str,
    *,
    required: bool = True,
) -> Any:
    target_key = aliases.get(canonical_key, canonical_key)
    if target_key in available:
        return available[target_key]
    if required:
        available_keys = ", ".join(sorted(available.keys())) or "none"
        raise ValueError(
            f"Missing required component input '{canonical_key}' (resolved as '{target_key}'). "
            f"Available inputs: {available_keys}."
        )
    return None
 
 
def _component_params(context: Any, process_name: str) -> dict[str, Any]:
    if context and hasattr(context, "get_parameters"):
        params = context.get_parameters(process_name)
        if isinstance(params, dict):
            return dict(params)
    return {}
 
 
def _filtered_component_params(params: dict[str, Any]) -> dict[str, Any]:
    return {k: v for k, v in params.items() if k not in _RESERVED_PARAM_KEYS}
 
 
def _run_csv_to_xy_component(context: Any, params: dict[str, Any]) -> dict[str, Any]:
    """
    Load a CSV file into X and y using project config and process parameters.

    Resolution rules:
    - CSV path:
      1) parameters.path or parameters.csv_path (when provided), interpreted
         as relative to the project root when not absolute
      2) global_config['data']['path'] from the project config
    - Label column:
      - parameters.label_column (required)
    """
    ctx = context or get_current_context()

    from pathlib import Path
    import os

    try:
        workspace_raw = os.environ.get("MLOPS_WORKSPACE_DIR") or "."
        workspace_root = Path(workspace_raw).expanduser().resolve()
    except Exception:
        workspace_root = Path.cwd()

    project_id = None
    try:
        project_id = getattr(ctx, "project_id", None)
    except Exception:
        project_id = None
    if not project_id:
        project_id = os.environ.get("MLOPS_PROJECT_ID")
    project_root = workspace_root / project_id if project_id else workspace_root

    csv_path_param = None
    try:
        if isinstance(params.get("path"), str) and params["path"].strip():
            csv_path_param = params["path"].strip()
        elif isinstance(params.get("csv_path"), str) and params["csv_path"].strip():
            csv_path_param = params["csv_path"].strip()
    except Exception:
        csv_path_param = None

    csv_rel: Path | None = None
    if csv_path_param:
        csv_rel = Path(csv_path_param)
    else:
        cfg = {}
        try:
            cfg_val = getattr(ctx, "global_config", {}) if ctx is not None else {}
            if isinstance(cfg_val, dict):
                cfg = cfg_val
        except Exception:
            cfg = {}

        data_cfg = cfg.get("data") if isinstance(cfg, dict) else None
        if not isinstance(data_cfg, dict):
            raise ValueError(
                "CSV data component requires a dataset path. "
                "Set 'data.path' in your project config or provide a 'path' parameter."
            )
        raw_path = data_cfg.get("path")
        if not isinstance(raw_path, str) or not raw_path.strip():
            raise ValueError(
                "Project config 'data.path' must be a non-empty string for the CSV data component."
            )
        csv_rel = Path(raw_path.strip())

    assert csv_rel is not None
    csv_path = csv_rel if csv_rel.is_absolute() else (project_root / csv_rel)

    if not csv_path.exists():
        raise FileNotFoundError(
            f"Training CSV not found at {csv_path}. "
            f"Expected the dataset at {csv_rel} relative to the project root."
        )

    import pandas as pd

    df = pd.read_csv(csv_path)

    label_column = params.get("label_column")
    if not isinstance(label_column, str) or not label_column.strip():
        available = ", ".join(map(str, df.columns)) or "<none>"
        raise ValueError(
            "CSV data component requires a non-empty 'label_column' parameter. "
            f"Available columns: {available}."
        )
    label_column = label_column.strip()

    if label_column not in df.columns:
        available = ", ".join(map(str, df.columns)) or "<none>"
        raise ValueError(
            f"Target column '{label_column}' not found in the CSV. "
            f"Available columns: {available}."
        )

    y = df.pop(label_column).values
    X = df.values

    return {"X": X, "y": y}
 
 
def _iterative_config_from_proc_cfg(proc_cfg: dict[str, Any]) -> dict[str, Any]:
    raw = proc_cfg.get("iterative_logging")
    if not isinstance(raw, dict):
        return {"enabled": False}

    metric_name = str(raw.get("metric_name", "train_loss")).strip() or "train_loss"
    eval_split = str(raw.get("eval_split", "train")).strip().lower() or "train"
    return {
        "enabled": True,
        "metric": metric_name,
        "eval_split": eval_split,
    }
 
 
def _slice_like(data: Any, start: int, end: int) -> Any:
    if hasattr(data, "iloc"):
        return data.iloc[start:end]
    return data[start:end]
 
 
def _resolve_eval_split_data(
    available_inputs: dict[str, Any],
    input_aliases: dict[str, str],
    split_name: str,
    train_X: Any,
    train_y: Any,
) -> tuple[Any, Any]:
    if split_name != "val":
        return train_X, train_y

    X_val = _resolve_input(available_inputs, input_aliases, "X_val", required=False)
    y_val = _resolve_input(available_inputs, input_aliases, "y_val", required=False)
    if X_val is None:
        X_val = _resolve_input(available_inputs, input_aliases, "X_test", required=False)
    if y_val is None:
        y_val = _resolve_input(available_inputs, input_aliases, "y_test", required=False)
    if X_val is None or y_val is None:
        logger.warning(
            "iterative_eval_split=val requested but val/test data not found; falling back to train split."
        )
        return train_X, train_y
    return X_val, y_val
 
 
def _compute_iterative_metric(metric_name: str, y_true: Any, y_pred: Any, model: Any) -> float:
    from sklearn.metrics import accuracy_score, mean_absolute_error, mean_squared_error, r2_score

    metric_id = str(metric_name).strip().lower()
    if metric_id in {"mse", "mean_squared_error"}:
        return float(mean_squared_error(y_true, y_pred))
    if metric_id in {"mae", "mean_absolute_error"}:
        return float(mean_absolute_error(y_true, y_pred))
    if metric_id in {"r2", "r2_score"}:
        return float(r2_score(y_true, y_pred))
    if "accuracy" in metric_id:
        return float(accuracy_score(y_true, y_pred))

    est_type = str(getattr(model, "_estimator_type", "")).strip().lower()
    if est_type == "regressor":
        return float(mean_squared_error(y_true, y_pred))
    return float(accuracy_score(y_true, y_pred))
 
 
def _fit_via_partial_fit(
    component: Any,
    X: Any,
    y: Any,
    metric_name: str,
    X_eval: Any,
    y_eval: Any,
    *,
    params: dict[str, Any],
) -> Any:
    model_cls = getattr(component.__class__, "_sklearn_class", None)
    if model_cls is None or not hasattr(model_cls, "partial_fit"):
        raise ValueError(
            "Iterative logging is only supported for sklearn components whose underlying "
            "estimator exposes a 'partial_fit' method."
        )
    if y is None:
        raise ValueError("Iterative logging with partial_fit requires target 'y', but no 'y' input was provided.")

    try:
        total_rows = len(X)
    except Exception as exc:  # pragma: no cover
        raise ValueError("Iterative logging requires a sized 'X' input.") from exc
    if total_rows <= 0:
        raise ValueError("Iterative logging requires at least one training sample.")

    try:
        batch_size = max(1, min(256, int(total_rows)))
    except Exception:  # pragma: no cover
        batch_size = 1

    try:
        model = model_cls(**component.params)
    except Exception:
        raise ValueError("Failed to construct estimator for iterative logging; check component hyperparameters.")
    if not hasattr(model, "partial_fit"):
        raise ValueError(
            "Iterative logging is only supported for sklearn components whose underlying "
            "estimator exposes a 'partial_fit' method."
        )

    classes = None
    try:
        import numpy as np

        classes = np.unique(y)
    except Exception:
        classes = None

    num_iterations = 1
    max_iter_param = None
    try:
        max_iter_param = int(params.get("max_iter")) if "max_iter" in params else None
    except Exception:
        max_iter_param = None
    max_iter_attr = None
    try:
        max_iter_attr = int(getattr(model, "max_iter"))
    except Exception:
        max_iter_attr = None

    if max_iter_param is not None and max_iter_param > 0:
        num_iterations = max_iter_param
    elif max_iter_attr is not None and max_iter_attr > 0:
        num_iterations = max_iter_attr

    metric_id = str(metric_name).strip().lower()

    for iteration in range(1, num_iterations + 1):
        start_idx = 0
        while start_idx < total_rows:
            end_idx = min(total_rows, start_idx + batch_size)
            X_batch = _slice_like(X, start_idx, end_idx)
            y_batch = _slice_like(y, start_idx, end_idx)
            call_kwargs: dict[str, Any] = {}
            if classes is not None and iteration == 1 and start_idx == 0:
                call_kwargs["classes"] = classes
            try:
                model.partial_fit(X_batch, y_batch, **call_kwargs)
            except TypeError:
                model.partial_fit(X_batch, y_batch)
            start_idx = end_idx

        try:
            if "loss" in metric_id and hasattr(model, "loss_"):
                log_metric(metric_name, float(model.loss_), step=iteration)
            elif X_eval is not None and y_eval is not None:
                y_pred = model.predict(X_eval)
                metric_val = _compute_iterative_metric(metric_name, y_eval, y_pred, model)
                log_metric(metric_name, float(metric_val), step=iteration)
        except Exception:
            continue

    component._fitted_model = model
    return component


def _fit_with_iterative_logging(
    adapter_cls: type,
    call_params: dict[str, Any],
    available_inputs: dict[str, Any],
    input_aliases: dict[str, str],
    X: Any,
    y: Any,
    iterative_cfg: dict[str, Any],
) -> Any:
    component = adapter_cls(**call_params)
    if not iterative_cfg.get("enabled"):
        component.fit(X, y)
        return component

    metric_name = str(iterative_cfg.get("metric", "train_loss")).strip() or "train_loss"
    X_eval, y_eval = _resolve_eval_split_data(
        available_inputs,
        input_aliases,
        str(iterative_cfg.get("eval_split", "train")),
        X,
        y,
    )

    return _fit_via_partial_fit(
        component=component,
        X=X,
        y=y,
        metric_name=metric_name,
        X_eval=X_eval,
        y_eval=y_eval,
        params=call_params,
    )


def _run_train_test_split_component(
    util_fn: Any,
    available_inputs: dict[str, Any],
    input_aliases: dict[str, str],
    params: dict[str, Any],
) -> dict[str, Any]:
    X = _resolve_input(available_inputs, input_aliases, "X", required=True)
    y = _resolve_input(available_inputs, input_aliases, "y", required=False)
    call_params = _filtered_component_params(params)

    if y is None:
        X_train, X_test = util_fn(X, **call_params)
        return {"X_train": X_train, "X_test": X_test}

    X_train, X_test, y_train, y_test = util_fn(X, y, **call_params)
    return {
        "X_train": X_train,
        "X_test": X_test,
        "y_train": y_train,
        "y_test": y_test,
    }


def _run_estimator_component(
    adapter_cls: type,
    available_inputs: dict[str, Any],
    input_aliases: dict[str, str],
    params: dict[str, Any],
    iterative_cfg: dict[str, Any] | None = None,
) -> dict[str, Any]:
    operation = str(params.get("operation", "fit")).strip().lower()
    call_params = _filtered_component_params(params)
    if iterative_cfg is None:
        iterative_cfg = {"enabled": False}

    if operation == "fit":
        X = _resolve_input(available_inputs, input_aliases, "X", required=True)
        y = _resolve_input(available_inputs, input_aliases, "y", required=False)
        component = _fit_with_iterative_logging(
            adapter_cls=adapter_cls,
            call_params=call_params,
            available_inputs=available_inputs,
            input_aliases=input_aliases,
            X=X,
            y=y,
            iterative_cfg=iterative_cfg,
        )
        return {"model": component._fitted_model}

    if operation == "fit_transform":
        X = _resolve_input(available_inputs, input_aliases, "X", required=True)
        y = _resolve_input(available_inputs, input_aliases, "y", required=False)
        component = adapter_cls(**call_params)
        X_transformed = component.fit_transform(X, y)
        return {"model": component._fitted_model, "X_transformed": X_transformed}

    if operation == "fit_predict":
        X = _resolve_input(available_inputs, input_aliases, "X", required=True)
        y = _resolve_input(available_inputs, input_aliases, "y", required=False)
        component = _fit_with_iterative_logging(
            adapter_cls=adapter_cls,
            call_params=call_params,
            available_inputs=available_inputs,
            input_aliases=input_aliases,
            X=X,
            y=y,
            iterative_cfg=iterative_cfg,
        )
        if not hasattr(component, "predict"):
            raise ValueError("Estimator does not support predict for operation 'fit_predict'.")
        predictions = component.predict(X)
        return {"model": component._fitted_model, "predictions": predictions}

    model = _resolve_input(available_inputs, input_aliases, "model", required=True)
    X = _resolve_input(available_inputs, input_aliases, "X", required=True)

    if operation == "transform":
        if not hasattr(model, "transform"):
            raise ValueError("Provided model does not support transform.")
        return {"X_transformed": model.transform(X)}

    if operation == "predict":
        if not hasattr(model, "predict"):
            raise ValueError("Provided model does not support predict.")
        return {"predictions": model.predict(X)}

    if operation == "predict_proba":
        if not hasattr(model, "predict_proba"):
            raise ValueError("Provided model does not support predict_proba.")
        return {"probabilities": model.predict_proba(X)}

    raise ValueError(
        f"Unsupported component operation '{operation}'. "
        "Supported operations: fit, fit_transform, fit_predict, transform, predict, predict_proba."
    )


@process(description="Execute a built-in sklearn component process")
def run_sklearn_component(context: Any = None, **kwargs: Any) -> dict[str, Any]:
    process_name = get_current_process_context() or getattr(context, "current_process", None) or "unknown_process"
    proc_cfg = _lookup_process_cfg(context, str(process_name))

    component_ref = str(proc_cfg.get("component", "")).strip()
    if not component_ref:
        raise ValueError(
            f"Process '{process_name}' is configured to use component runner but has no 'component' field."
        )
    if not component_ref.startswith(COMPONENT_PREFIX):
        raise ValueError(
            f"Unsupported component '{component_ref}'. "
            f"Component references must start with '{COMPONENT_PREFIX}'."
        )

    legacy_keys = [k for k in ("inputs", "outputs") if k in proc_cfg]
    if legacy_keys:
        joined = ", ".join(sorted(legacy_keys))
        raise ValueError(
            f"Process '{process_name}' is configured as a sklearn component but uses legacy key(s) "
            f"{joined}. Use 'input_transform' for input aliasing and 'output_mapping' for output "
            "aliasing."
        )

    component_body = component_ref[len(COMPONENT_PREFIX) :].strip()
    if not component_body:
        raise ValueError(f"Invalid component reference '{component_ref}'.")

    component_name = component_body
    operation_suffix: str | None = None
    if "." in component_body:
        base, op = component_body.rsplit(".", 1)
        base = base.strip()
        op = op.strip()
        if base and op:
            component_name = base
            operation_suffix = op

    input_aliases = _alias_map(proc_cfg, "input_transform")
    output_aliases = _alias_map(proc_cfg, "output_mapping")
    params = _component_params(context, str(process_name))
    iterative_cfg = _iterative_config_from_proc_cfg(proc_cfg)
    available_inputs = dict(kwargs)

    if component_name == "csv_to_xy":
        if operation_suffix is not None:
            raise ValueError(
                f"Utility-like component '{component_ref}' must not specify an operation suffix. "
                "Use a plain 'sklearn.csv_to_xy' reference."
            )
        result = _run_csv_to_xy_component(context, params)
        return {output_aliases.get(k, k): v for k, v in result.items()}

    raw_params = proc_cfg.get("parameters")
    if isinstance(raw_params, dict) and "operation" in raw_params:
        raise ValueError(
            f"Process '{process_name}' is configured with 'parameters.operation', which is no longer "
            "supported for sklearn components. Encode the operation in the 'component' string instead, "
            "for example: component: 'sklearn.LogisticRegression.fit'."
        )

    utility = get_utility(component_name)
    if utility is not None:
        if operation_suffix is not None:
            raise ValueError(
                f"Utility sklearn component '{component_ref}' must not specify an operation suffix. "
                "Use a plain 'sklearn.<function_name>' reference (for example 'sklearn.train_test_split')."
            )
        result = _run_train_test_split_component(
            utility.func,
            available_inputs,
            input_aliases,
            params,
        )
        return {output_aliases.get(k, k): v for k, v in result.items()}

    if operation_suffix is None:
        raise ValueError(
            f"Estimator sklearn component '{component_ref}' must specify an explicit operation method. "
            "Use the form 'sklearn.<EstimatorClass>.<operation>' (for example "
            "'sklearn.LogisticRegression.fit' or 'sklearn.MLPClassifier.predict')."
        )
    params["operation"] = operation_suffix

    adapter_cls = get_adapter(component_name)
    if adapter_cls is None:
        raise ValueError(
            f"Unknown sklearn component '{component_ref}'. "
            "Check component name and installed sklearn version."
        )
    result = _run_estimator_component(
        adapter_cls,
        available_inputs,
        input_aliases,
        params,
        iterative_cfg=iterative_cfg,
    )
    return {output_aliases.get(k, k): v for k, v in result.items()}


def ensure_component_runner_registered() -> None:
    """Ensure component runner remains registered even after test/global resets."""
    process_registry = get_process_registry()
    if process_registry and process_registry.get_process(COMPONENT_RUNNER_NAME) is None:
        process_def = ProcessDefinition(
            name=COMPONENT_RUNNER_NAME,
            description="Execute a built-in sklearn component process",
            parameters={},
            step_names=[],
            runner=run_sklearn_component,
            logging=True,
        )
        process_def.original_func = run_sklearn_component
        process_registry.register_process(process_def)

