package store

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"

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

type FrozenWorkflow struct {
	Name       string
	Definition api.WorkflowDefinition
	Brief      string
	Title      string
}

func (s *Store) ListWorkflows(ctx context.Context) ([]api.WorkflowSummary, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT name, coalesce(json_extract(definition, '$.description'), '') FROM workflows ORDER BY name`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	workflows := []api.WorkflowSummary{}
	for rows.Next() {
		w := api.WorkflowSummary{}
		err = rows.Scan(&w.Name, &w.Description)
		if err != nil {
			return nil, err
		}
		workflows = append(workflows, w)
	}
	return workflows, rows.Err()
}

func (s *Store) GetWorkflow(ctx context.Context, name string) (api.WorkflowDefinition, error) {
	def := ""
	err := s.db.QueryRowContext(ctx, `SELECT definition FROM workflows WHERE name = ?`, name).Scan(&def)
	if errors.Is(err, sql.ErrNoRows) {
		return api.WorkflowDefinition{}, fmt.Errorf("%w: no workflow %q", ErrNotFound, name)
	}
	if err != nil {
		return api.WorkflowDefinition{}, err
	}
	w := api.WorkflowDefinition{}
	err = json.Unmarshal([]byte(def), &w)
	return w, err
}

var ErrWorkflowExists = errors.New("workflow exists")

func (s *Store) PutWorkflow(ctx context.Context, name, from string, def api.WorkflowDefinition, actor string) error {
	b, err := json.Marshal(def)
	if err != nil {
		return err
	}
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		if from != "" {
			found := false
			err := tx.QueryRowContext(ctx, `SELECT EXISTS (SELECT 1 FROM workflows WHERE name = ?)`, from).Scan(&found)
			if err != nil {
				return err
			}
			if !found {
				return fmt.Errorf("%w: no workflow %q", ErrNotFound, from)
			}
		}
		if from != "" && from != name {
			taken := false
			err := tx.QueryRowContext(ctx, `SELECT EXISTS (SELECT 1 FROM workflows WHERE name = ?)`, name).Scan(&taken)
			if err != nil {
				return err
			}
			if taken {
				return fmt.Errorf("%w: a workflow named %q already exists", ErrWorkflowExists, name)
			}
			_, err = tx.ExecContext(ctx, `DELETE FROM workflows WHERE name = ?`, from)
			if err != nil {
				return err
			}
			err = repointReviews(ctx, tx, emit, from, name, actor)
			if err != nil {
				return err
			}
			err = emit(api.Event{Kind: "workflow.deleted", Actor: actor, Payload: map[string]any{"name": from}})
			if err != nil {
				return err
			}
		}
		_, err := tx.ExecContext(ctx, `INSERT INTO workflows (name, definition) VALUES (?, ?)
			ON CONFLICT (name) DO UPDATE SET definition = excluded.definition`, name, string(b))
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "workflow.updated", Actor: actor, Payload: map[string]any{"name": name}})
	})
}

func repointReviews(ctx context.Context, tx *sql.Tx, emit emitFunc, from, to, actor string) error {
	rows, err := tx.QueryContext(ctx, `UPDATE projects SET review_workflow = ? WHERE review_workflow = ? RETURNING `+projectColumns, to, from)
	if err != nil {
		return err
	}
	projects := []api.Project{}
	for rows.Next() {
		p, err := scanProject(rows)
		if err != nil {
			rows.Close()
			return err
		}
		projects = append(projects, p)
	}
	err = errors.Join(rows.Err(), rows.Close())
	if err != nil {
		return err
	}
	for _, p := range projects {
		err = emit(api.Event{Kind: "project.updated", ProjectId: &p.Id, Actor: actor, Payload: ToPayload(p)})
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) DeleteWorkflow(ctx context.Context, name, actor string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		res, err := tx.ExecContext(ctx, `DELETE FROM workflows WHERE name = ?`, name)
		if err != nil {
			return err
		}
		n, err := res.RowsAffected()
		if err != nil {
			return err
		}
		if n == 0 {
			return fmt.Errorf("%w: no workflow %q", ErrNotFound, name)
		}
		err = repointReviews(ctx, tx, emit, name, "", actor)
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "workflow.deleted", Actor: actor, Payload: map[string]any{"name": name}})
	})
}

func (s *Store) FreezeWorkflow(ctx context.Context, name string) (FrozenWorkflow, error) {
	def, err := s.GetWorkflow(ctx, name)
	if err != nil {
		return FrozenWorkflow{}, err
	}
	return FrozenWorkflow{Name: name, Definition: def}, nil
}

func (s *Store) TaskWorkflow(ctx context.Context, taskID string) (FrozenWorkflow, error) {
	f := FrozenWorkflow{}
	def := ""
	err := s.db.QueryRowContext(ctx, `SELECT workflow_name, workflow_definition, workflow_brief, workflow_title FROM tasks WHERE id = ?`, taskID).Scan(&f.Name, &def, &f.Brief, &f.Title)
	if errors.Is(err, sql.ErrNoRows) {
		return FrozenWorkflow{}, fmt.Errorf("%w: no task %s", ErrNotFound, taskID)
	}
	if err != nil {
		return FrozenWorkflow{}, err
	}
	if f.Name == "" {
		return FrozenWorkflow{}, fmt.Errorf("%w: task %s has no workflow", ErrNotFound, taskID)
	}
	err = json.Unmarshal([]byte(def), &f.Definition)
	return f, err
}

func (s *Store) RunTasks(ctx context.Context, run string) ([]api.Task, error) {
	return s.queryTasks(ctx, `SELECT `+taskColumns+` FROM tasks WHERE workflow_run = ? ORDER BY created_at`, run)
}

func (s *Store) RunEvents(ctx context.Context, run string) ([]api.Event, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT `+eventColumns+` FROM events
		WHERE task_id IN (SELECT id FROM tasks WHERE workflow_run = ?) AND kind IN ('task.workflow_stage', 'task.discarded', 'task.workflow_finished')
		ORDER BY seq
	`, run)
	if err != nil {
		return nil, err
	}
	return scanEvents(rows)
}

func (s *Store) RecordStage(ctx context.Context, taskID, output, commit string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE tasks SET stage_output = ?, stage_commit = ? WHERE id = ?`, output, commit, taskID)
		return err
	})
}

func (s *Store) StageResult(ctx context.Context, taskID string) (string, string, error) {
	output := ""
	commit := ""
	err := s.db.QueryRowContext(ctx, `SELECT stage_output, stage_commit FROM tasks WHERE id = ?`, taskID).Scan(&output, &commit)
	return output, commit, err
}

func (s *Store) SetStageOutput(ctx context.Context, sessionID, output string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = ? WHERE id = ?`, output, sessionID)
		return err
	})
}

func (s *Store) FinishStage(ctx context.Context, sessionID string) (bool, error) {
	finished := false
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		res, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_finished = 1 WHERE id = ? AND stage_output != '' AND NOT EXISTS (SELECT 1 FROM tasks WHERE tasks.id = agent_sessions.task_id AND tasks.discard_state != '')`, sessionID)
		if err != nil {
			return err
		}
		n, err := res.RowsAffected()
		finished = n > 0
		return err
	})
	return finished, err
}

func (s *Store) DropStageOutput(ctx context.Context, sessionID string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = '' WHERE id = ? AND stage_finished = 0`, sessionID)
		return err
	})
}

func (s *Store) StageOutput(ctx context.Context, sessionID string) (string, error) {
	output := ""
	err := s.db.QueryRowContext(ctx, `SELECT stage_output FROM agent_sessions WHERE id = ?`, sessionID).Scan(&output)
	return output, err
}

func (s *Store) FinishedStages(ctx context.Context) ([]string, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT id FROM agent_sessions WHERE stage_finished = 1 ORDER BY created_at`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	ids := []string{}
	for rows.Next() {
		id := ""
		err = rows.Scan(&id)
		if err != nil {
			return nil, err
		}
		ids = append(ids, id)
	}
	return ids, rows.Err()
}

func (s *Store) ClearStage(ctx context.Context, sessionID string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		return clearStage(ctx, tx, sessionID)
	})
}

func clearStage(ctx context.Context, tx *sql.Tx, sessionID string) error {
	_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = '', stage_finished = 0 WHERE id = ?`, sessionID)
	return err
}

func (s *Store) FinishWorkflow(ctx context.Context, task api.Task, sessionID, workflowName, output string, kept []string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		err := clearStage(ctx, tx, sessionID)
		if err != nil {
			return err
		}
		return emit(api.Event{
			Kind:      "task.workflow_finished",
			ProjectId: &task.ProjectId,
			TaskId:    &task.Id,
			Actor:     "byre",
			Payload:   map[string]any{"id": task.Id, "workflow": workflowName, "output": output, "kept": kept},
		})
	})
}

func (s *Store) WorkflowFinished(ctx context.Context, run string) (bool, error) {
	finished := false
	err := s.db.QueryRowContext(ctx, `SELECT EXISTS (SELECT 1 FROM events e JOIN tasks t ON t.id = e.task_id WHERE t.workflow_run = ? AND e.kind = 'task.workflow_finished')`, run).Scan(&finished)
	return finished, err
}