package main

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
)

type direction string

const (
	dirUp   direction = "up"
	dirDown direction = "down"
)

type migrationFile struct {
	Version    uint
	Identifier string
	Direction  direction
	Filename   string // e.g. "0000001_INITIAL_CREATE.up.sql"
}

type migrationSource struct {
	dir        string
	index      []uint                                // sorted unique version numbers
	migrations map[uint]map[direction]*migrationFile // version -> direction -> file
}

// matches: 0000001_INITIAL_CREATE.up.sql
var migrationRegex = regexp.MustCompile(`^([0-9]+)_(.*)\.(up|down)\.(.*)$`)

func newMigrationSource(dir string) (*migrationSource, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, fmt.Errorf("failed to read migrations directory: %w", err)
	}

	s := &migrationSource{
		dir:        dir,
		migrations: make(map[uint]map[direction]*migrationFile),
	}

	versionSet := make(map[uint]struct{})

	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}

		name := entry.Name()
		m := migrationRegex.FindStringSubmatch(name)
		if len(m) != 5 {
			// skip non-matching filenames (same as library)
			continue
		}

		versionUint64, err := strconv.ParseUint(m[1], 10, 64)
		if err != nil {
			continue
		}

		version := uint(versionUint64)
		dir := direction(m[3])

		mf := &migrationFile{
			Version:    version,
			Identifier: m[2],
			Direction:  dir,
			Filename:   name,
		}

		if _, ok := s.migrations[version]; !ok {
			s.migrations[version] = make(map[direction]*migrationFile)
		}

		if _, exists := s.migrations[version][dir]; exists {
			return nil, fmt.Errorf("duplicate migration version %d direction %s", version, dir)
		}

		s.migrations[version][dir] = mf
		versionSet[version] = struct{}{}
	}

	// build sorted index
	s.index = make([]uint, 0, len(versionSet))
	for v := range versionSet {
		s.index = append(s.index, v)
	}
	sort.Slice(s.index, func(i, j int) bool {
		return s.index[i] < s.index[j]
	})

	return s, nil
}

func (s *migrationSource) first() (uint, error) {
	if len(s.index) == 0 {
		return 0, os.ErrNotExist
	}
	return s.index[0], nil
}

func (s *migrationSource) next(version uint) (uint, error) {
	pos := sort.Search(len(s.index), func(i int) bool {
		return s.index[i] > version
	})
	if pos >= len(s.index) {
		return 0, os.ErrNotExist
	}
	return s.index[pos], nil
}

func (s *migrationSource) prev(version uint) (uint, error) {
	pos := sort.Search(len(s.index), func(i int) bool {
		return s.index[i] >= version
	})
	// pos is the index of version (or where it would be inserted)
	// we want the one before it
	if pos <= 0 {
		return 0, os.ErrNotExist
	}
	return s.index[pos-1], nil
}

func (s *migrationSource) readUp(version uint) (string, error) {
	if dirs, ok := s.migrations[version]; ok {
		if mf, ok := dirs[dirUp]; ok {
			return filepath.Join(s.dir, mf.Filename), nil
		}
	}
	return "", os.ErrNotExist
}

func (s *migrationSource) readDown(version uint) (string, error) {
	if dirs, ok := s.migrations[version]; ok {
		if mf, ok := dirs[dirDown]; ok {
			return filepath.Join(s.dir, mf.Filename), nil
		}
	}
	return "", os.ErrNotExist
}

func (s *migrationSource) versionExists(version uint) bool {
	pos := sort.Search(len(s.index), func(i int) bool {
		return s.index[i] >= version
	})
	return pos < len(s.index) && s.index[pos] == version
}