package store

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"maps"
	"slices"

	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/workflow"
)

type oldStage struct {
	Name   string            `json:"name"`
	After  *[]string         `json:"after"`
	Base   *string           `json:"base"`
	Model  string            `json:"model"`
	Prompt string            `json:"prompt"`
	Agents *[]api.TeamMember `json:"agents"`
}

type oldDefinition struct {
	Description string           `json:"description"`
	Stages      []oldStage       `json:"stages"`
	Steps       *json.RawMessage `json:"steps"`
}

func migrateWorkflowSteps(tx *sql.Tx) error {
	err := convertColumn(tx, "workflows", "name", "definition")
	if err != nil {
		return err
	}
	return convertColumn(tx, "tasks", "id", "workflow_definition")
}

func convertColumn(tx *sql.Tx, table, key, column string) error {
	rows, err := tx.Query(fmt.Sprintf(`SELECT %s, %s FROM %s WHERE %s != ''`, key, column, table, column))
	if err != nil {
		return err
	}
	type row struct{ key, def string }
	all := []row{}
	for rows.Next() {
		r := row{}
		err = rows.Scan(&r.key, &r.def)
		if err != nil {
			rows.Close()
			return err
		}
		all = append(all, r)
	}
	rows.Close()
	err = rows.Err()
	if err != nil {
		return err
	}

	for _, r := range all {
		old := oldDefinition{}
		err = json.Unmarshal([]byte(r.def), &old)
		if err != nil {
			return fmt.Errorf("%s %s: %w", table, r.key, err)
		}
		if old.Steps != nil {
			continue
		}
		def, err := toSteps(old)
		if err != nil {
			return fmt.Errorf("%s %s: %w", table, r.key, err)
		}
		j, err := json.Marshal(def)
		if err != nil {
			return err
		}
		_, err = tx.Exec(fmt.Sprintf(`UPDATE %s SET %s = ? WHERE %s = ?`, table, column, key), string(j), r.key)
		if err != nil {
			return err
		}
	}
	return nil
}

func toSteps(old oldDefinition) (api.WorkflowDefinition, error) {
	after := map[string][]string{}
	byName := map[string]oldStage{}
	for i, s := range old.Stages {
		if s.Base != nil && *s.Base != "parent" {
			return api.WorkflowDefinition{}, fmt.Errorf("stage %q has base %q, which steps can't express", s.Name, *s.Base)
		}
		byName[s.Name] = s
		switch {
		case s.After != nil:
			after[s.Name] = *s.After
		case i == 0:
			after[s.Name] = []string{}
		default:
			after[s.Name] = []string{old.Stages[i-1].Name}
		}
	}

	pending := func(names []string, done map[string]bool) []string {
		out := []string{}
		for _, n := range names {
			if !done[n] {
				out = append(out, n)
			}
		}
		return out
	}
	var closure func(name string, done map[string]bool) map[string]bool
	closure = func(name string, done map[string]bool) map[string]bool {
		c := map[string]bool{name: true}
		for _, a := range pending(after[name], done) {
			for n := range closure(a, done) {
				c[n] = true
			}
		}
		return c
	}
	sequence := func(steps ...api.WorkflowStep) api.WorkflowStep {
		order := []api.WorkflowStep{}
		for _, s := range steps {
			if s.Order != nil {
				order = append(order, *s.Order...)
				continue
			}
			order = append(order, s)
		}
		return api.WorkflowStep{Order: &order}
	}

	var tree func(ends []string, done map[string]bool) (api.WorkflowStep, error)
	tree = func(ends []string, done map[string]bool) (api.WorkflowStep, error) {
		groups := [][]string{}
		for _, e := range ends {
			merged := []string{e}
			kept := [][]string{}
			for _, g := range groups {
				if slices.ContainsFunc(g, func(other string) bool { return shareAny(closure(e, done), closure(other, done)) }) {
					merged = append(g, merged...)
					continue
				}
				kept = append(kept, g)
			}
			groups = append(kept, merged)
		}
		if len(groups) > 1 {
			branches := []api.WorkflowStep{}
			for _, g := range groups {
				b, err := tree(g, done)
				if err != nil {
					return api.WorkflowStep{}, err
				}
				branches = append(branches, b)
			}
			return api.WorkflowStep{Parallel: &branches}, nil
		}

		if len(ends) > 1 {
			common := closure(ends[0], done)
			for _, e := range ends[1:] {
				c := closure(e, done)
				for n := range common {
					if !c[n] {
						delete(common, n)
					}
				}
			}
			heads := []string{}
			for n := range common {
				later := false
				for m := range common {
					later = later || (m != n && closure(m, done)[n])
				}
				if !later {
					heads = append(heads, n)
				}
			}
			if len(heads) == 0 {
				return api.WorkflowStep{}, fmt.Errorf("stages %v overlap in earlier stages without sharing all of them, which steps can't express", ends)
			}
			slices.Sort(heads)
			prefix, err := tree(heads, done)
			if err != nil {
				return api.WorkflowStep{}, err
			}
			next := maps.Clone(done)
			for n := range common {
				next[n] = true
			}
			rest, err := tree(ends, next)
			if err != nil {
				return api.WorkflowStep{}, err
			}
			return sequence(prefix, rest), nil
		}

		s := byName[ends[0]]
		stage := api.WorkflowStep{Stage: s.Name, Model: s.Model, Prompt: s.Prompt, Agents: s.Agents}
		preds := pending(after[s.Name], done)
		if len(preds) == 0 {
			return stage, nil
		}
		before, err := tree(preds, done)
		if err != nil {
			return api.WorkflowStep{}, err
		}
		return sequence(before, stage), nil
	}

	finals := []string{}
	for _, s := range old.Stages {
		next := false
		for _, a := range after {
			next = next || slices.Contains(a, s.Name)
		}
		if !next {
			finals = append(finals, s.Name)
		}
	}
	top, err := tree(finals, map[string]bool{})
	if err != nil {
		return api.WorkflowDefinition{}, err
	}
	steps := []api.WorkflowStep{top}
	if top.Order != nil {
		steps = *top.Order
	}
	def := api.WorkflowDefinition{Description: old.Description, Steps: steps}

	if len(workflow.Stages(def)) != len(old.Stages) {
		return api.WorkflowDefinition{}, fmt.Errorf("converted to %d stages from %d", len(workflow.Stages(def)), len(old.Stages))
	}
	for _, s := range old.Stages {
		got := workflow.After(def, s.Name)
		want := after[s.Name]
		if !slices.Equal(sorted(got), sorted(want)) {
			return api.WorkflowDefinition{}, fmt.Errorf("stage %q converted to come after %v, not %v", s.Name, got, want)
		}
	}
	return def, nil
}

func shareAny(a, b map[string]bool) bool {
	for n := range a {
		if b[n] {
			return true
		}
	}
	return false
}

func sorted(names []string) []string {
	out := slices.Clone(names)
	slices.Sort(out)
	return out
}