# pipelines/training_pipeline.py
from kfp import dsl
from kfp.components import load_component_from_file


@dsl.component(base_image="python:3.10-slim")
def preprocess_data(input_path: str, output_path: str):
    import pandas as pd
    from sklearn.model_selection import train_test_split
    from sklearn.preprocessing import StandardScaler
    import joblib
    import os

    df = pd.read_csv(input_path)
    scaler = StandardScaler()
    feature_cols = [c for c in df.columns if c != "target"]
    df[feature_cols] = scaler.fit_transform(df[feature_cols])
    train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)

    os.makedirs(output_path, exist_ok=True)
    train_df.to_csv(f"{output_path}/train.csv", index=False)
    test_df.to_csv(f"{output_path}/test.csv", index=False)
    joblib.dump(scaler, f"{output_path}/scaler.joblib")


@dsl.component(base_image="python:3.10-slim")
def train_model(data_path: str, model_path: str, params: dict):
    import pandas as pd
    import lightgbm as lgb
    import mlflow
    import json
    import os

    train_df = pd.read_csv(f"{data_path}/train.csv")
    X_train = train_df.drop("target", axis=1)
    y_train = train_df["target"]

    mlflow.set_experiment("production-training")
    with mlflow.start_run():
        mlflow.log_params(params)

        train_data = lgb.Dataset(X_train, label=y_train)
        model = lgb.train(
            params, train_data, num_boost_round=500,
            valid_sets=[train_data], callbacks=[lgb.log_evaluation(50)],
        )

        mlflow.lightgbm.log_model(model, "model")
        os.makedirs(model_path, exist_ok=True)
        model.save_model(f"{model_path}/model.txt")
        mlflow.log_artifact(f"{model_path}/model.txt")


@dsl.component(base_image="python:3.10-slim")
def evaluate_model(data_path: str, model_path: str, metrics_path: str):
    import pandas as pd
    import lightgbm as lgb
    from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
    import json
    import os

    test_df = pd.read_csv(f"{data_path}/test.csv")
    X_test = test_df.drop("target", axis=1)
    y_test = test_df["target"]

    model = lgb.Booster(model_file=f"{model_path}/model.txt")
    y_pred = (model.predict(X_test) > 0.5).astype(int)

    metrics = {
        "accuracy": float(accuracy_score(y_test, y_pred)),
        "f1": float(f1_score(y_test, y_pred)),
        "auc": float(roc_auc_score(y_test, model.predict(X_test))),
    }

    os.makedirs(metrics_path, exist_ok=True)
    with open(f"{metrics_path}/metrics.json", "w") as f:
        json.dump(metrics, f, indent=2)
    print(f"Metrics: {metrics}")


@dsl.component(base_image="python:3.10-slim")
def validate_model(metrics_path: str, threshold: float = 0.85):
    import json
    with open(f"{metrics_path}/metrics.json") as f:
        metrics = json.load(f)
    assert metrics["accuracy"] >= threshold, f"Model accuracy {metrics['accuracy']} below threshold {threshold}"
    print(f"Model validation passed: accuracy={metrics['accuracy']:.4f}")


@dsl.pipeline(
    name="training-pipeline",
    description="End-to-end ML training pipeline",
)
def training_pipeline(
    input_path: str = "gs://ml-data/input/",
    output_path: str = "gs://ml-data/processed/",
    model_path: str = "gs://ml-models/production/",
    metrics_path: str = "gs://ml-metrics/",
):
    preprocess_task = preprocess_data(input_path=input_path, output_path=output_path)

    train_task = train_model(
        data_path=preprocess_task.outputs["output_path"],
        model_path=model_path,
        params={"num_leaves": 31, "learning_rate": 0.05, "objective": "binary"},
    )

    eval_task = evaluate_model(
        data_path=preprocess_task.outputs["output_path"],
        model_path=train_task.outputs["model_path"],
        metrics_path=metrics_path,
    )

    validate_task = validate_model(metrics_path=eval_task.outputs["metrics_path"])

    # Set dependency
    validate_task.after(train_task)

# seldon/model.py
import numpy as np
import lightgbm as lgb
import logging
from typing import Dict, List

logger = logging.getLogger(__name__)


class SeldonModel:
    def __init__(self):
        self.model = None
        self.scaler = None

    def load(self, model_path: str):
        self.model = lgb.Booster(model_file=f"{model_path}/model.txt")
        logger.info(f"Loaded model from {model_path}")

    def predict(self, X: np.ndarray, features_names: List[str] = None) -> np.ndarray:
        """Seldon predict interface."""
        predictions = self.model.predict(X)
        return (predictions > 0.5).astype(int).reshape(-1, 1)

    def predict_proba(self, X: np.ndarray, features_names: List[str] = None) -> np.ndarray:
        """Return prediction probabilities."""
        proba = self.model.predict(X)
        return np.column_stack([1 - proba, proba])

    def feedback(self, feedback_request) -> bool:
        """Handle feedback for online learning."""
        logger.info(f"Received feedback: {feedback_request}")
        return True

# feature_store/feature_definitions.py
from feast import FeatureStore, Entity, Feature, ValueType
from feast import FileSource, FeatureView
from datetime import timedelta
import pandas as pd


# Define entities
user_entity = Entity(name="user_id", value_type=ValueType.INT64, description="User ID")
product_entity = Entity(name="product_id", value_type=ValueType.INT64, description="Product ID")


# Define feature views
user_features = FeatureView(
    name="user_features",
    entities=["user_id"],
    ttl=timedelta(days=1),
    features=[
        Feature(name="purchase_count_7d", value_type=ValueType.INT64),
        Feature(name="avg_order_value", value_type=ValueType.FLOAT),
        Feature(name="last_purchase_days_ago", value_type=ValueType.INT64),
    ],
    online=True,
    source=FileSource(path="data/user_features.parquet", timestamp_field="event_timestamp"),
)

product_features = FeatureView(
    name="product_features",
    entities=["product_id"],
    ttl=timedelta(days=7),
    features=[
        Feature(name="price", value_type=ValueType.FLOAT),
        Feature(name="category_id", value_type=ValueType.INT64),
        Feature(name="avg_rating", value_type=ValueType.FLOAT),
        Feature(name="review_count", value_type=ValueType.INT64),
    ],
    online=True,
    source=FileSource(path="data/product_features.parquet", timestamp_field="event_timestamp"),
)


# Feature retrieval
store = FeatureStore(repo_path=".")

training_df = store.get_historical_features(
    entity_df=entity_df,
    features=["user_features:purchase_count_7d", "product_features:price"],
).to_df()

online_features = store.get_online_features(
    features=["user_features:purchase_count_7d", "product_features:price"],
    entity_rows=[{"user_id": 1, "product_id": 42}],
).to_dict()