"""
FoundryMatch SaaS - ML Model Training Pipeline
===============================================
This module provides a training pipeline for the Matching Intelligence
engine. It uses user feedback to train a model that can predict match
confidence with higher accuracy than simple scoring.
"""

import json
import logging
import pickle
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

import asyncpg
import numpy as np
import pandas as pd
import mlflow
import mlflow.sklearn

# Use Agg backend for headless servers
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

# Try to import seaborn, fall back if not available
try:
    import seaborn as sns

    SEABORN_AVAILABLE = True
except ImportError:
    SEABORN_AVAILABLE = False
from rapidfuzz import fuzz
from sklearn.ensemble import RandomForestClassifier
from sklearn.calibration import CalibratedClassifierCV
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, accuracy_score, confusion_matrix

log = logging.getLogger(__name__)


@dataclass
class ModelArtifact:
    """Represents a trained model and its metadata."""

    model: RandomForestClassifier
    model_version: str
    training_date: datetime
    metrics: Dict[str, Any]
    feature_names: List[str]
    mappings_used: List[Dict[str, Any]]


class MatchingIntelligenceTrainer:
    """
    Trains a machine learning model to predict match confidence based on
    user feedback and detailed feature analysis.
    """

    def __init__(self, db_pool: asyncpg.Pool):
        """
        Initializes the trainer with a database connection pool.

        Args:
            db_pool: An active asyncpg connection pool.
        """
        self.db_pool = db_pool
        # Define the feature calculation functions
        self.feature_functions = {
            "wratio": fuzz.WRatio,
            "partial_ratio": fuzz.partial_ratio,
            "token_set_ratio": fuzz.token_set_ratio,
            "token_sort_ratio": fuzz.token_sort_ratio,
            "jaro_winkler": fuzz.jaro_winkler_similarity,
        }

    async def fetch_training_data(
        self, db=None, lookback_days: int = 90
    ) -> Optional[pd.DataFrame]:
        """
        Fetches user feedback and the associated record data for training.
        Joins feedback with historical match results.

        Args:
            db: Database connection (optional, uses self.db_pool if not provided)
            lookback_days: How many days of feedback to use for training.

        Returns:
            A DataFrame containing labeled pairs of records for training.
        """
        log.info(f"Fetching training data from the last {lookback_days} days...")

        # Real implementation: join feedback with match results
        query = """
            SELECT
                r.field,
                r.source_value,
                r.ref_value,
                r.features_json,
                CASE
                    WHEN f.label = 'accept' THEN 1
                    WHEN f.label = 'reject' THEN 0
                END as label
            FROM match_results r
            JOIN feedback f ON f.job_id = r.job_id AND f.pair_id = r.pair_id
            WHERE f.created_at >= NOW() - $1 * INTERVAL '1 day'
              AND f.label IN ('accept', 'reject')
            LIMIT 10000;
        """

        try:
            conn = db or self.db_pool
            data = await conn.fetch(query, lookback_days)
            if not data:
                raise ValueError("No training data found in the database.")

            df = pd.DataFrame(data)

            # Parse features_json if it exists
            if "features_json" in df.columns and df["features_json"].notna().any():
                df["features"] = df.apply(
                    lambda row: json.loads(row["features_json"])
                    if row["features_json"]
                    else self._compute_features(row["source_value"], row["ref_value"]),
                    axis=1,
                )
            else:
                # Compute features on the fly from RapidFuzz
                df["features"] = df.apply(
                    lambda row: self._compute_features(
                        row["source_value"], row["ref_value"]
                    ),
                    axis=1,
                )

            log.info(f"Fetched {len(df)} labeled pairs from the database.")
            return df

        except Exception as e:
            log.warning(
                f"Could not fetch real training data ({e}). Using enhanced dummy data."
            )
            # TODO: Remove this once tables are properly set up
            dummy_data = []
            for _ in range(1000):
                source_val = f"Company_{np.random.randint(1, 100)}"
                ref_val = f"Company_{np.random.randint(1, 100)}"
                label = 1 if np.random.random() > 0.3 else 0
                dummy_data.append(
                    {
                        "field": "name",
                        "source_value": source_val,
                        "ref_value": ref_val,
                        "features": self._compute_features(source_val, ref_val),
                        "label": label,
                    }
                )
            return pd.DataFrame(dummy_data)

    def _compute_features(self, source_val: str, ref_val: str) -> Dict[str, float]:
        """Compute features from RapidFuzz for a pair of values."""
        if not source_val or not ref_val:
            return {name: 0.0 for name in self.feature_functions.keys()}

        features = {}
        for name, func in self.feature_functions.items():
            try:
                score = func(str(source_val), str(ref_val))
                # Handle both 0-100 and 0-1 score ranges safely
                features[name] = score / 100.0 if score > 1 else float(score)
            except:
                features[name] = 0.0
        return features

    def _generate_features_for_pair(
        self,
        source_record: Dict[str, Any],
        ref_record: Dict[str, Any],
        mappings: List[Dict[str, str]],
    ) -> Dict[str, float]:
        """Generates a flat feature dictionary for a single pair of records."""
        features = {}
        for mapping in mappings:
            s_col, r_col = mapping["source"], mapping["ref"]
            s_val = str(source_record.get(s_col, "")).strip()
            r_val = str(ref_record.get(r_col, "")).strip()

            # Calculate all features for this field, defaulting to 0.0 if data is missing
            for func_name, func in self.feature_functions.items():
                feature_name = f"{s_col}_{func_name}"
                if s_val and r_val:
                    score = func(s_val, r_val)
                    # Handle both 0-100 and 0-1 score ranges safely
                    features[feature_name] = (
                        score / 100.0 if score > 1 else float(score)
                    )
                else:
                    features[feature_name] = 0.0

        return features

    def train(
        self, training_df: pd.DataFrame, model_params: Optional[Dict[str, Any]] = None
    ) -> Optional[ModelArtifact]:
        """
        Trains a RandomForestClassifier with calibration on the provided data.

        Args:
            training_df: DataFrame with 'features' and 'label' columns.
            model_params: Optional parameters for the RandomForestClassifier.

        Returns:
            A ModelArtifact containing the trained model and metadata, or None.
        """
        if training_df.empty:
            log.error("Training data is empty.")
            return None

        log.info(f"Starting model training on {len(training_df)} samples.")

        # Start MLflow run
        with mlflow.start_run():
            # 1. Feature Engineering
            if "features" in training_df.columns:
                X = pd.DataFrame(training_df["features"].tolist())
            else:
                # Fallback to old format
                mappings = training_df.get(
                    "mappings", [{"source": "field", "ref": "field"}]
                ).iloc[0]
                features_list = training_df.apply(
                    lambda row: self._generate_features_for_pair(
                        row.get(
                            "source_record", {"field": row.get("source_value", "")}
                        ),
                        row.get(
                            "reference_record", {"field": row.get("ref_value", "")}
                        ),
                        mappings,
                    ),
                    axis=1,
                )
                X = pd.DataFrame(features_list.tolist())

            feature_names = list(X.columns)

            # 2. Target Variable
            if "label" in training_df.columns:
                y = training_df["label"]
            else:
                y = training_df["status"].apply(lambda s: 1 if s == "approved" else 0)

            # 3. Train/Test Split
            X_train, X_test, y_train, y_test = train_test_split(
                X,
                y,
                test_size=0.25,
                random_state=42,
                stratify=y if len(np.unique(y)) > 1 else None,
            )

            # 4. Model Training - Baseline RandomForest
            default_params = {
                "n_estimators": 150,
                "random_state": 42,
                "n_jobs": -1,
                "max_depth": 10,
            }
            if model_params:
                default_params.update(model_params)

            base_model = RandomForestClassifier(**default_params)
            log.info(f"Training RandomForestClassifier with params: {default_params}")
            base_model.fit(X_train, y_train)

            # 5. Calibration for better confidence scores
            model = CalibratedClassifierCV(base_model, method="sigmoid", cv=3)
            model.fit(X_train, y_train)

            # 6. Evaluation
            y_pred = model.predict(X_test)
            y_proba = model.predict_proba(X_test)[:, 1]
            accuracy = accuracy_score(y_test, y_pred)
            report = classification_report(y_test, y_pred, output_dict=True)
            cm = confusion_matrix(y_test, y_pred)

            log.info(f"Model training complete. Accuracy: {accuracy:.2%}")
            log.info(f"Classification Report:\n{classification_report(y_test, y_pred)}")

            # Log to MLflow
            mlflow.log_params(default_params)
            mlflow.log_metrics(
                {
                    "accuracy": accuracy,
                    "precision": report["weighted avg"]["precision"],
                    "recall": report["weighted avg"]["recall"],
                    "f1_score": report["weighted avg"]["f1-score"],
                }
            )

            # Create and log confusion matrix
            fig, ax = plt.subplots(figsize=(8, 6))
            if SEABORN_AVAILABLE:
                sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", ax=ax)
            else:
                # Fallback to matplotlib
                im = ax.imshow(cm, cmap="Blues")
                plt.colorbar(im, ax=ax)
                # Add text annotations
                for i in range(cm.shape[0]):
                    for j in range(cm.shape[1]):
                        ax.text(j, i, str(cm[i, j]), ha="center", va="center")
            ax.set_xlabel("Predicted")
            ax.set_ylabel("Actual")
            ax.set_title("Confusion Matrix")
            mlflow.log_figure(fig, "confusion_matrix.png")
            plt.close(fig)

            # Log model
            mlflow.sklearn.log_model(
                model, "model", registered_model_name="forgeI_mapping_scorer"
            )

            metrics = {
                "accuracy": accuracy,
                "classification_report": report,
                "confusion_matrix": cm.tolist(),
                "feature_importances": dict(
                    zip(feature_names, base_model.feature_importances_)
                ),
            }

            # 7. Create Artifact
            artifact = ModelArtifact(
                model=model,
                model_version=f"v{datetime.utcnow().strftime('%Y%m%d%H%M%S')}",
                training_date=datetime.utcnow(),
                metrics=metrics,
                feature_names=feature_names,
                mappings_used=training_df.get("mappings", []).iloc[0]
                if "mappings" in training_df.columns
                else [],
            )

            return artifact

    def save_model(self, artifact: ModelArtifact, path: Path) -> None:
        """Saves the model artifact to a file using pickle."""
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "wb") as f:
            pickle.dump(artifact, f)
        log.info(f"Model artifact '{artifact.model_version}' saved to {path}")

    async def run_training_pipeline(
        self, lookback_days: int = 90, model_output_path: Optional[Path] = None
    ) -> Optional[Path]:
        """Orchestrates the full training pipeline."""
        log.info("===== Starting Matching Intelligence Training Pipeline =====")

        training_df = await self.fetch_training_data(lookback_days=lookback_days)
        if training_df is None or training_df.empty:
            log.error("Training pipeline aborted: No data.")
            return None

        artifact = self.train(training_df)
        if artifact is None:
            log.error("Training pipeline aborted: Model training failed.")
            return None

        if model_output_path:
            self.save_model(artifact, model_output_path)
            log.info(
                f"===== Pipeline Finished Successfully. Model saved to {model_output_path} ====="
            )
            return model_output_path
        else:
            log.info(
                "===== Pipeline Finished Successfully. Model not saved (no path provided). ====="
            )
            return None
