package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/jackc/pgx/v5"
)

type User struct {
	ID       uint64 `json:"id"`
	Username string `json:"username"`
	Roles    []Role `json:"roles"`
}

type Role struct {
	ID       uint64 `json:"id"`
	Rolename string `json:"rolename"`
}

func main() {
	ctx := context.Background()
	conn, err := pgx.Connect(ctx, os.Getenv("DSN"))
	if err != nil {
		log.Fatalln(err)
	}
	defer conn.Close(ctx)

	users := []User{}

	t := time.Now()
	for i := 0; i < 1000; i++ {
		users, err = listUsers(ctx, conn)
		if err != nil {
			log.Fatalln(err)
		}
	}

	enc := json.NewEncoder(os.Stdout)
	enc.SetIndent("", "\t")

	if err := enc.Encode(users); err != nil {
		log.Fatalln(err)
	}

	fmt.Println("Duration", time.Since(t).String())

	t = time.Now()
	for i := 0; i < 1000; i++ {
		users, err = listUsersOptimized(ctx, conn)
		if err != nil {
			log.Fatalln(err)
		}

	}

	if err := enc.Encode(users); err != nil {
		log.Fatalln(err)
	}

	fmt.Println("Duration", time.Since(t).String())
}

func listUsers(ctx context.Context, conn *pgx.Conn) ([]User, error) {
	rows, err := conn.Query(ctx, "SELECT id, username FROM users")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	users := []User{}
	for rows.Next() {
		u := new(User)
		if err := rows.Scan(&u.ID, &u.Username); err != nil {
			return nil, err
		}
		users = append(users, *u)
	}

	rows, err = conn.Query(ctx, "SELECT id, rolename FROM roles")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	roles := []Role{}
	for rows.Next() {
		r := new(Role)
		if err := rows.Scan(&r.ID, &r.Rolename); err != nil {
			return nil, err
		}
		roles = append(roles, *r)
	}

	rows, err = conn.Query(ctx, "SELECT user_id, role_id FROM users_roles")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var userID, roleID uint64
		if err := rows.Scan(&userID, &roleID); err != nil {
			return nil, err
		}
		for i := range users {
			if users[i].ID == userID {
				for j := range roles {
					if roles[j].ID == roleID {
						users[i].Roles = append(users[i].Roles, roles[j])
					}
				}
			}
		}
	}

	return users, nil
}

func listUsersOptimized(ctx context.Context, conn *pgx.Conn) ([]User, error) {
	rows, err := conn.Query(ctx, `
		SELECT id, username, (
			SELECT JSONB_AGG(r.*)
			FROM (
				SELECT id, rolename
				FROM roles INNER JOIN users_roles AS ur ON id = ur.role_id AND u.id = ur.user_id
			) AS r
		) AS roles
		FROM users AS u
		GROUP BY id, username`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return pgx.CollectRows(rows, pgx.RowToStructByPos[User])
}