from sqlmodel import Session, select
from models import User, Post
from typing import List
from collections import defaultdict
from models import UserWithRecentPosts


def get_users_with_recent_posts(session: Session) -> List[UserWithRecentPosts]:
    users = session.exec(select(User)).all()
    user_ids = [u.id for u in users]

    posts = session.exec(
        select(Post)
        .where(Post.user_id.in_(user_ids))
        .order_by(Post.user_id, Post.timestamp.desc())
    ).all()

    # Group posts by user_id
    posts_by_user = defaultdict(list)
    for post in posts:
        if len(posts_by_user[post.user_id]) < 5:
            posts_by_user[post.user_id].append(post)

    return [
        UserWithRecentPosts(
            id=user.id,
            username=user.username,
            display_name=user.display_name,
            recent_posts=posts_by_user.get(user.id, []),
        )
        for user in users
    ]