from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import numpy as np
from app.services.database_service import db
from app.utils.logger import debug_logger


def collect_interaction_data():
    """
    Collect user interaction data for training the prediction models.
    """
    interactions = list(
        db.interactions.find(
            {},
            {
                "user_id": 1,
                "item_id": 1,
                "interaction_type": 1,  # e.g., like, comment, follow
                "timestamp": 1,
            },
        )
    )
    return interactions


def create_features(interactions):
    """
    Create features for users, posts, and stories based on interaction data.
    """
    user_features = {}
    post_features = {}
    story_features = {}

    for interaction in interactions:
        user_id = interaction["user_id"]
        item_id = interaction["item_id"]
        interaction_type = interaction["interaction_type"]

        # Increment feature counts based on interaction type
        if interaction_type == "like":
            user_features.setdefault(user_id, {}).setdefault("likes", 0)
            user_features[user_id]["likes"] += 1
        elif interaction_type == "comment":
            user_features.setdefault(user_id, {}).setdefault("comments", 0)
            user_features[user_id]["comments"] += 1
        elif interaction_type == "follow":
            user_features.setdefault(user_id, {}).setdefault("follows", 0)
            user_features[user_id]["follows"] += 1

        # Update post and story features
        if item_id.startswith("post_"):
            post_features.setdefault(item_id, {}).setdefault("engagement", 0)
            post_features[item_id]["engagement"] += 1
        elif item_id.startswith("story_"):
            story_features.setdefault(item_id, {}).setdefault("views", 0)
            story_features[item_id]["views"] += 1

    return user_features, post_features, story_features


def train_model(features, labels):
    """
    Train a machine learning model using the provided features and labels.
    """
    X_train, X_test, y_train, y_test = train_test_split(
        features, labels, test_size=0.2, random_state=42
    )
    model = RandomForestClassifier()
    model.fit(X_train, y_train)
    accuracy = model.score(X_test, y_test)
    debug_logger.info(f"Model trained with accuracy: {accuracy}")
    return model


def update_model(model, new_interactions):
    """
    Update the model with new interaction data.
    """
    new_features = [interaction["features"] for interaction in new_interactions]
    new_labels = [
        1 if interaction["interaction_type"] in ["follow", "like", "view"] else 0
        for interaction in new_interactions
    ]
    model.partial_fit(new_features, new_labels)
    return model


def recommend_followers(user_id, model, n_recs=10):
    """
    Recommend followers for a user using the trained model.
    """
    user_features = get_user_features(user_id)
    predictions = model.predict_proba([user_features])[:, 1]  # Probability of follow
    recommended_users = np.argsort(predictions)[-n_recs:]
    return recommended_users


def recommend_posts(user_id, model, n_recs=10):
    """
    Recommend posts for a user using the trained model.
    """
    user_features = get_user_features(user_id)
    post_features = get_post_features()
    predictions = model.predict_proba(post_features)[:, 1]  # Probability of engagement
    recommended_posts = np.argsort(predictions)[-n_recs:]
    return recommended_posts


def recommend_stories(user_id, model, n_recs=10):
    """
    Recommend stories for a user using the trained model.
    """
    user_features = get_user_features(user_id)
    story_features = get_story_features()
    predictions = model.predict_proba(story_features)[:, 1]  # Probability of engagement
    recommended_stories = np.argsort(predictions)[-n_recs:]
    return recommended_stories


# Helper functions to fetch features
def get_user_features(user_id):
    """
    Fetch features for a specific user.
    """
    user = db.users.find_one({"user_id": user_id})
    if not user:
        return {}
    return {
        "likes": user.get("likes", 0),
        "comments": user.get("comments", 0),
        "follows": user.get("follows", 0),
    }


def get_post_features():
    """
    Fetch features for all posts.
    """
    posts = list(db.posts.find({}, {"post_id": 1, "engagement": 1}))
    return [post.get("engagement", 0) for post in posts]


def get_story_features():
    """
    Fetch features for all stories.
    """
    stories = list(db.stories.find({}, {"story_id": 1, "views": 1}))
    return [story.get("views", 0) for story in stories]