pulse-api/internal/repository/habit_freeze.go

package repository

import (
	"errors"
	"time"

	"github.com/daniil/homelab-api/internal/model"
	"github.com/jmoiron/sqlx"
)

var ErrFreezeNotFound = errors.New("freeze not found")
var ErrInvalidDateRange = errors.New("invalid date range")

type HabitFreezeRepository struct {
	db *sqlx.DB
}

func NewHabitFreezeRepository(db *sqlx.DB) *HabitFreezeRepository {
	return &HabitFreezeRepository{db: db}
}

func (r *HabitFreezeRepository) Create(freeze *model.HabitFreeze) error {
	// Validate date range
	if freeze.EndDate.Before(freeze.StartDate) {
		return ErrInvalidDateRange
	}

	query := `
		INSERT INTO habit_freezes (habit_id, user_id, start_date, end_date, reason)
		VALUES ($1, $2, $3, $4, $5)
		RETURNING id, created_at`

	return r.db.QueryRow(query,
		freeze.HabitID,
		freeze.UserID,
		freeze.StartDate,
		freeze.EndDate,
		freeze.Reason,
	).Scan(&freeze.ID, &freeze.CreatedAt)
}

func (r *HabitFreezeRepository) GetByHabitID(habitID, userID int64) ([]model.HabitFreeze, error) {
	query := `
		SELECT id, habit_id, user_id, start_date, end_date, reason, created_at
		FROM habit_freezes
		WHERE habit_id = $1 AND user_id = $2
		ORDER BY start_date DESC`

	var freezes []model.HabitFreeze
	if err := r.db.Select(&freezes, query, habitID, userID); err != nil {
		return nil, err
	}

	if freezes == nil {
		freezes = []model.HabitFreeze{}
	}

	return freezes, nil
}

func (r *HabitFreezeRepository) GetActiveForHabit(habitID int64, date time.Time) (*model.HabitFreeze, error) {
	query := `
		SELECT id, habit_id, user_id, start_date, end_date, reason, created_at
		FROM habit_freezes
		WHERE habit_id = $1 AND start_date <= $2 AND end_date >= $2`

	var freeze model.HabitFreeze
	err := r.db.Get(&freeze, query, habitID, date)
	if err != nil {
		return nil, err
	}

	return &freeze, nil
}

func (r *HabitFreezeRepository) IsHabitFrozenOnDate(habitID int64, date time.Time) (bool, error) {
	query := `
		SELECT COUNT(*) FROM habit_freezes
		WHERE habit_id = $1 AND start_date <= $2 AND end_date >= $2`

	var count int
	err := r.db.Get(&count, query, habitID, date)
	if err != nil {
		return false, err
	}

	return count > 0, nil
}

func (r *HabitFreezeRepository) GetFreezesForDateRange(habitID int64, startDate, endDate time.Time) ([]model.HabitFreeze, error) {
	query := `
		SELECT id, habit_id, user_id, start_date, end_date, reason, created_at
		FROM habit_freezes
		WHERE habit_id = $1 
		AND NOT (end_date < $2 OR start_date > $3)
		ORDER BY start_date`

	var freezes []model.HabitFreeze
	if err := r.db.Select(&freezes, query, habitID, startDate, endDate); err != nil {
		return nil, err
	}

	if freezes == nil {
		freezes = []model.HabitFreeze{}
	}

	return freezes, nil
}

func (r *HabitFreezeRepository) Delete(freezeID, userID int64) error {
	query := `DELETE FROM habit_freezes WHERE id = $1 AND user_id = $2`
	result, err := r.db.Exec(query, freezeID, userID)
	if err != nil {
		return err
	}

	rows, _ := result.RowsAffected()
	if rows == 0 {
		return ErrFreezeNotFound
	}

	return nil
}

func (r *HabitFreezeRepository) CountFrozenDaysInRange(habitID int64, startDate, endDate time.Time) (int, error) {
	freezes, err := r.GetFreezesForDateRange(habitID, startDate, endDate)
	if err != nil {
		return 0, err
	}

	frozenDays := 0
	for _, freeze := range freezes {
		// Calculate overlap between freeze period and query range
		overlapStart := freeze.StartDate
		if startDate.After(freeze.StartDate) {
			overlapStart = startDate
		}
		
		overlapEnd := freeze.EndDate
		if endDate.Before(freeze.EndDate) {
			overlapEnd = endDate
		}

		if !overlapEnd.Before(overlapStart) {
			days := int(overlapEnd.Sub(overlapStart).Hours()/24) + 1
			frozenDays += days
		}
	}

	return frozenDays, nil
}