"""Utilities for connecting to and interacting with Hopsworks Feature Store."""

import os
import logging
from typing import Optional
import pandas as pd
import hopsworks
from dotenv import load_dotenv

logger = logging.getLogger(__name__)

# Load environment variables
load_dotenv()


def connect_to_hopsworks(api_key: Optional[str] = None, project_name: Optional[str] = None):
    """
    Connect to Hopsworks.

    Args:
        api_key: Hopsworks API key. If None, reads from HOPSWORKS_API_KEY env variable.
        project_name: Hopsworks project name. If None, reads from HOPSWORKS_PROJECT_NAME env variable.

    Returns:
        Hopsworks project object

    Raises:
        ValueError: If API key or project name is not provided
    """
    # Get API key from parameter or environment
    api_key = api_key or os.getenv("HOPSWORKS_API_KEY")
    if not api_key:
        raise ValueError(
            "Hopsworks API key not provided. "
            "Set HOPSWORKS_API_KEY environment variable or pass api_key parameter."
        )

    # Get project name from parameter or environment
    project_name = project_name or os.getenv("HOPSWORKS_PROJECT")

    logger.info(f"Connecting to Hopsworks project: {project_name or 'default'}")

    # Login to Hopsworks
    try:
        if project_name:
            project = hopsworks.login(
                api_key_value=api_key,
                project=project_name,
                engine="python"  # Use Python engine (serverless, no cert download)
            )
        else:
            project = hopsworks.login(
                api_key_value=api_key,
                engine="python"
            )
    except Exception as e:
        logger.error(f"Failed to connect to Hopsworks: {e}")
        logger.info("Trying to connect without specifying project...")
        project = hopsworks.login(
            api_key_value=api_key,
            engine="python"
        )

    logger.info(f"Successfully connected to Hopsworks project: {project.name}")
    return project


def get_or_create_feature_group(
    project,
    name: str,
    version: int = 1,
):
    """
    Get existing feature group or create new one if it doesn't exist.

    Args:
        project: Hopsworks project object
        name: Feature group name
        version: Feature group version
        description: Description of the feature group
        primary_key: List of column names to use as primary key
        event_time: Column name to use as event time
        online_enabled: Whether to enable online feature serving

    Returns:
        Feature group object
    """
    fs = project.get_feature_store()

    try:
        # Try to get existing feature group
        fg = fs.get_feature_group(name=name, version=version)
        logger.info(f"Retrieved existing feature group: {name} (version {version})")
        return fg
    except Exception:
        # Feature group doesn't exist, will need to create it
        logger.info(f"Feature group {name} (version {version}) not found, will create on first insert")
        return None


def upload_dataframe_to_feature_group(
    project,
    df: pd.DataFrame,
    feature_group_name: str,
    version: int = 1,
    description: str = "",
    primary_key: list = None,
    event_time: Optional[str] = None,
    online_enabled: bool = False,
    write_options: dict = None
):
    """
    Upload a DataFrame to a Hopsworks feature group.

    Args:
        project: Hopsworks project object
        df: Pandas DataFrame to upload
        feature_group_name: Name of the feature group
        version: Feature group version
        description: Description of the feature group
        primary_key: List of column names to use as primary key
        event_time: Column name to use as event time
        online_enabled: Whether to enable online feature serving
        write_options: Additional write options (e.g., {"wait_for_job": False})

    Returns:
        Feature group object
    """
    fs = project.get_feature_store()

    logger.info(f"Uploading DataFrame to feature group: {feature_group_name} (version {version})")
    logger.info(f"DataFrame shape: {df.shape}")

    # Create or get feature group
    fg = fs.get_or_create_feature_group(
        name=feature_group_name,
        version=version,
        description=description,
        primary_key=primary_key or [],
        event_time=event_time,
        online_enabled=online_enabled
    )

    # Insert data
    write_options = write_options or {"wait_for_job": True}
    fg.insert(df, write_options=write_options)

    logger.info(f"Successfully uploaded {len(df)} rows to {feature_group_name}")
    return fg


def read_feature_group(
    project,
    feature_group_name: str,
    version: int = 1,
    online: bool = False
) -> pd.DataFrame:
    """
    Read data from a Hopsworks feature group.

    Args:
        project: Hopsworks project object
        feature_group_name: Name of the feature group
        version: Feature group version
        online: Whether to read from online feature store

    Returns:
        Pandas DataFrame with feature group data
    """
    fs = project.get_feature_store()

    logger.info(f"Reading feature group: {feature_group_name} (version {version})")

    fg = fs.get_feature_group(name=feature_group_name, version=version)

    if online:
        df = fg.read(online=True)
    else:
        df = fg.read()

    logger.info(f"Read {len(df)} rows from {feature_group_name}")
    return df


def create_feature_view(
    project,
    name: str,
    version: int = 1,
    description: str = "",
    query=None,
    labels: list = None
):
    """
    Create a feature view for training datasets.

    Args:
        project: Hopsworks project object
        name: Feature view name
        version: Feature view version
        description: Description of the feature view
        query: Query object to define feature selection
        labels: List of label column names

    Returns:
        Feature view object
    """
    fs = project.get_feature_store()

    logger.info(f"Creating feature view: {name} (version {version})")

    fv = fs.create_feature_view(
        name=name,
        version=version,
        description=description,
        query=query,
        labels=labels or []
    )

    logger.info(f"Successfully created feature view: {name}")
    return fv


def save_model_to_registry(
    project,
    model,
    model_name: str,
    metrics: dict = None,
    description: str = "",
    model_schema: dict = None,
    scaler=None,
    feature_names: list = None
):
    """
    Save a trained model to Hopsworks Model Registry with all artifacts.

    Args:
        project: Hopsworks project object
        model: Trained model object
        model_name: Name for the model in registry
        metrics: Dictionary of model metrics
        description: Model description
        model_schema: Optional model schema
        scaler: Optional scaler object to save with model
        feature_names: Optional list of feature names

    Returns:
        Model registry object
    """
    import joblib
    import os
    import tempfile

    logger.info(f"Saving model to Hopsworks Model Registry: {model_name}")

    # Get model registry
    mr = project.get_model_registry()

    # Create temporary directory for model artifacts
    with tempfile.TemporaryDirectory() as tmpdir:
        # Save model using joblib
        model_path = os.path.join(tmpdir, "model.pkl")
        joblib.dump(model, model_path)
        logger.info(f"Model saved to temporary path: {model_path}")

        # Save scaler if provided
        if scaler is not None:
            scaler_path = os.path.join(tmpdir, "scaler.pkl")
            joblib.dump(scaler, scaler_path)
            logger.info(f"Scaler saved to temporary path: {scaler_path}")

        # Save feature names if provided
        if feature_names is not None:
            feature_names_path = os.path.join(tmpdir, "feature_names.txt")
            with open(feature_names_path, 'w') as f:
                f.write('\n'.join(feature_names))
            logger.info(f"Feature names saved to temporary path: {feature_names_path}")

        # Save hyperparameters from description if they exist
        if description and "Best Parameters:" in description:
            import json
            # Try to extract and save params as JSON for easier loading
            params_path = os.path.join(tmpdir, "hyperparameters.txt")
            with open(params_path, 'w') as f:
                f.write(description)
            logger.info(f"Hyperparameters saved to temporary path: {params_path}")

        # Create model in registry
        try:
            model_registry = mr.python.create_model(
                name=model_name,
                metrics=metrics or {},
                description=description,
                input_example=None,
                model_schema=model_schema
            )

            # Save all model artifacts (model, scaler, feature_names)
            model_registry.save(tmpdir)
            logger.info(f"Successfully saved model '{model_name}' with all artifacts to Model Registry")

            return model_registry

        except Exception as e:
            logger.error(f"Failed to save model to registry: {e}")
            raise