"""Orchestrator: generate_report() — fits metrics + saves all PNGs + JSON."""
import json
from pathlib import Path

import numpy as np

from ._metrics import (
    compute_classification_metrics,
    compute_regression_metrics,
    compute_georf_metrics,
)
from ._plots import (
    plot_confusion_matrix,
    plot_roc_curves,
    plot_feature_importances,
    plot_actual_vs_predicted,
    plot_residuals,
    plot_residual_distribution,
    plot_tree_quality_scores,
)


def _sniff_task(estimator):
    """Return 'classification' or 'regression' by duck-typing the estimator."""
    if hasattr(estimator, "predict_proba"):
        return "classification"
    return "regression"


def generate_report(
    estimator,
    X_test,
    y_test,
    task="auto",
    output_dir=".",
    prefix="report",
    feature_names=None,
):
    """Generate a full metrics + visualization report for any sklearn estimator.

    Parameters
    ----------
    estimator : fitted sklearn-compatible estimator
    X_test : array-like of shape (n_samples, n_features)
    y_test : array-like of shape (n_samples,)
    task : 'auto' | 'classification' | 'regression'
        'auto' sniffs via hasattr(estimator, 'predict_proba').
    output_dir : str | Path
        Directory to write files into; created if missing.
    prefix : str
        Filename prefix, e.g. 'report' → report_metrics.json.
    feature_names : list[str] | None

    Returns
    -------
    dict[str, Path]
        Maps file role (e.g. 'metrics', 'confusion_matrix') to absolute Path.
    """
    out_dir = Path(output_dir).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    X_test = np.asarray(X_test)
    y_test = np.asarray(y_test)

    if task == "auto":
        task = _sniff_task(estimator)

    y_pred = estimator.predict(X_test)
    outputs = {}

    if task == "classification":
        outputs.update(_classification_report(
            estimator, X_test, y_test, y_pred,
            out_dir, prefix, feature_names,
        ))
    else:
        outputs.update(_regression_report(
            estimator, X_test, y_test, y_pred,
            out_dir, prefix, feature_names,
        ))

    return outputs


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _classification_report(estimator, X_test, y_test, y_pred,
                            out_dir, prefix, feature_names):
    outputs = {}

    # Probabilities (optional)
    y_proba = None
    if hasattr(estimator, "predict_proba"):
        try:
            y_proba = estimator.predict_proba(X_test)
        except Exception:
            pass

    # Classes
    classes = None
    if hasattr(estimator, "classes_"):
        classes = estimator.classes_

    # --- metrics JSON ---
    metrics = compute_classification_metrics(y_test, y_pred, y_proba=y_proba, classes=classes)
    georf = compute_georf_metrics(estimator, feature_names=feature_names)
    if georf:
        metrics["georf"] = georf

    json_path = out_dir / f"{prefix}_metrics.json"
    json_path.write_text(json.dumps(metrics, indent=2, default=str), encoding="utf-8")
    outputs["metrics"] = json_path

    # --- confusion matrix ---
    cm_path = out_dir / f"{prefix}_confusion_matrix.png"
    result = plot_confusion_matrix(y_test, y_pred, classes=classes, output_path=cm_path)
    if result:
        outputs["confusion_matrix"] = result

    # --- ROC curves ---
    if y_proba is not None:
        roc_path = out_dir / f"{prefix}_roc_curves.png"
        result = plot_roc_curves(y_test, y_proba, classes=classes, output_path=roc_path)
        if result:
            outputs["roc_curves"] = result

    # --- feature importances ---
    fi = _get_feature_importances(estimator, feature_names)
    if fi:
        fi_path = out_dir / f"{prefix}_feature_importances.png"
        result = plot_feature_importances(fi, output_path=fi_path)
        if result:
            outputs["feature_importances"] = result

    # --- tree quality scores (GeoRF) ---
    if hasattr(estimator, "tree_quality_scores"):
        try:
            scores = estimator.tree_quality_scores()
            tqs_path = out_dir / f"{prefix}_tree_quality_scores.png"
            result = plot_tree_quality_scores(scores, output_path=tqs_path)
            if result:
                outputs["tree_quality_scores"] = result
        except Exception:
            pass

    return outputs


def _regression_report(estimator, X_test, y_test, y_pred,
                        out_dir, prefix, feature_names):
    outputs = {}

    # --- metrics JSON ---
    metrics = compute_regression_metrics(y_test, y_pred)
    georf = compute_georf_metrics(estimator, feature_names=feature_names)
    if georf:
        metrics["georf"] = georf

    json_path = out_dir / f"{prefix}_metrics.json"
    json_path.write_text(json.dumps(metrics, indent=2, default=str), encoding="utf-8")
    outputs["metrics"] = json_path

    # --- actual vs predicted ---
    avp_path = out_dir / f"{prefix}_actual_vs_predicted.png"
    result = plot_actual_vs_predicted(y_test, y_pred, output_path=avp_path)
    if result:
        outputs["actual_vs_predicted"] = result

    # --- residuals scatter ---
    res_path = out_dir / f"{prefix}_residuals.png"
    result = plot_residuals(y_test, y_pred, output_path=res_path)
    if result:
        outputs["residuals"] = result

    # --- residual distribution ---
    rd_path = out_dir / f"{prefix}_residual_distribution.png"
    result = plot_residual_distribution(y_test, y_pred, output_path=rd_path)
    if result:
        outputs["residual_distribution"] = result

    # --- feature importances ---
    fi = _get_feature_importances(estimator, feature_names)
    if fi:
        fi_path = out_dir / f"{prefix}_feature_importances.png"
        result = plot_feature_importances(fi, output_path=fi_path)
        if result:
            outputs["feature_importances"] = result

    # --- tree quality scores (GeoRF) ---
    if hasattr(estimator, "tree_quality_scores"):
        try:
            scores = estimator.tree_quality_scores()
            tqs_path = out_dir / f"{prefix}_tree_quality_scores.png"
            result = plot_tree_quality_scores(scores, output_path=tqs_path)
            if result:
                outputs["tree_quality_scores"] = result
        except Exception:
            pass

    return outputs


def _get_feature_importances(estimator, feature_names):
    """Return {name: importance} dict from the estimator, or {} if not available."""
    if hasattr(estimator, "feature_importances"):
        try:
            return estimator.feature_importances(feature_names=feature_names)
        except Exception:
            pass
    if hasattr(estimator, "feature_importances_"):
        try:
            importances = np.asarray(estimator.feature_importances_)
            n = len(importances)
            names = feature_names if feature_names is not None else [f"feature_{i}" for i in range(n)]
            return dict(sorted(zip(names, importances.tolist()), key=lambda x: x[1], reverse=True))
        except Exception:
            pass
    return {}
