package services

import (
	"fmt"
	"os"
	"strings"
	"time"

	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/db"
	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/grpc/pb"
	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/models"
	"github.com/google/uuid"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

const stepDispatchGrace = 15 * time.Second

func TriggerWorkflow(instanceID, workflowID, actor string) (string, error) {
	wf, err := GetWorkflow(instanceID, workflowID)
	if err != nil {
		return "", err
	}
	targets, err := ResolveTargets(instanceID, wf.TargetServerIDs, wf.TargetTags)
	if err != nil {
		return "", err
	}
	if len(wf.Steps) == 0 {
		return "", fmt.Errorf("workflow has no steps")
	}

	ctx, cancel := wfCtx()
	running := db.Col("workflow_runs").FindOne(ctx, bson.M{"instance_id": instanceID, "workflow_id": workflowID, "status": "running"})
	cancel()
	if running.Err() == nil {
		return "", fmt.Errorf("workflow already has a run in progress")
	}

	resolved, err := resolveSteps(instanceID, wf)
	if err != nil {
		return "", err
	}

	run := models.WorkflowRun{
		InstanceID:  instanceID,
		RunID:       uuid.New().String(),
		WorkflowID:  workflowID,
		Name:        wf.Name,
		Steps:       resolved,
		Status:      "running",
		TriggeredBy: actor,
		StartedAt:   time.Now(),
		ServerRuns:  make([]models.ServerRun, 0, len(targets)),
	}
	for _, srv := range targets {
		sr := models.ServerRun{ServerID: srv.ServerID, Hostname: srv.Hostname, Status: "queued", RunEnv: map[string]string{}}
		for _, rs := range resolved {
			sr.Steps = append(sr.Steps, models.StepRun{Order: rs.Order, Name: rs.Name, Status: "queued", OutputEnv: map[string]string{}})
		}
		run.ServerRuns = append(run.ServerRuns, sr)
	}

	ictx, icancel := wfCtx()
	defer icancel()
	if _, err := db.Col("workflow_runs").InsertOne(ictx, run); err != nil {
		return "", err
	}

	go executeRun(run.RunID)
	return run.RunID, nil
}

func resolveSteps(instanceID string, wf *models.Workflow) ([]models.ResolvedStep, error) {
	ctx, cancel := wfCtx()
	defer cancel()
	out := make([]models.ResolvedStep, 0, len(wf.Steps))
	for _, ref := range wf.Steps {
		if ref.Inline != nil {
			out = append(out, resolveInlineStep(ref))
			continue
		}
		lib, err := getStep(ctx, instanceID, ref.StepID)
		if err != nil {
			return nil, err
		}
		inputs := map[string]string{}
		for _, p := range lib.DeclaredInputs {
			if ref.Inputs != nil {
				if v, ok := ref.Inputs[p.Name]; ok {
					inputs[p.Name] = v
					continue
				}
			}
			inputs[p.Name] = p.Default
		}
		rs := models.ResolvedStep{
			Order:       ref.Order,
			Name:        lib.Name,
			Interpreter: lib.Interpreter,
			Script:      lib.Script,
			SecretRefs:  lib.SecretRefs,
			OnFailure:   ref.OnFailure,
			MaxRetries:  ref.MaxRetries,
			Inputs:      inputs,
		}
		if ref.Overrides != nil {
			if ref.Overrides.Script != nil {
				rs.Script = *ref.Overrides.Script
			}
			if ref.Overrides.SecretRefs != nil {
				rs.SecretRefs = ref.Overrides.SecretRefs
			}
		}
		if rs.OnFailure == "" {
			rs.OnFailure = "stop"
		}
		out = append(out, rs)
	}
	return out, nil
}

func resolveInlineStep(ref models.WorkflowStepRef) models.ResolvedStep {
	in := ref.Inline
	inputs := map[string]string{}
	for _, p := range in.DeclaredInputs {
		if ref.Inputs != nil {
			if v, ok := ref.Inputs[p.Name]; ok {
				inputs[p.Name] = v
				continue
			}
		}
		inputs[p.Name] = p.Default
	}
	onFailure := ref.OnFailure
	if onFailure == "" {
		onFailure = "stop"
	}
	return models.ResolvedStep{
		Order:       ref.Order,
		Name:        in.Name,
		Interpreter: in.Interpreter,
		Script:      in.Script,
		SecretRefs:  in.SecretRefs,
		OnFailure:   onFailure,
		MaxRetries:  ref.MaxRetries,
		Inputs:      inputs,
	}
}

func executeRun(runID string) {
	run, err := getRunByID(runID)
	if err != nil {
		return
	}
	done := make(chan int, len(run.ServerRuns))
	for i := range run.ServerRuns {
		go func(idx int) {
			runServer(run.InstanceID, runID, idx, run.Steps, run.ServerRuns[idx].ServerID)
			done <- idx
		}(i)
	}
	for range run.ServerRuns {
		<-done
	}

	final, _ := getRunByID(runID)
	status := "success"
	for _, sr := range final.ServerRuns {
		if sr.Status == "failed" {
			status = "failed"
		}
	}
	now := time.Now()
	ctx, cancel := wfCtx()
	defer cancel()
	_, _ = db.Col("workflow_runs").UpdateOne(ctx, bson.M{"run_id": runID, "status": "running"},
		bson.M{"$set": bson.M{"status": status, "finished_at": now}})
}

func runServer(instanceID, runID string, srvIdx int, steps []models.ResolvedStep, serverID string) {
	now := time.Now()
	setServerRun(runID, srvIdx, bson.M{"server_runs.$.status": "running", "server_runs.$.started_at": now})

	if !Dispatcher.IsConnected(serverID) {
		fin := time.Now()
		_, _ = AppendMarker(runID, serverID, "agent not connected server skipped")
		setServerRun(runID, srvIdx, bson.M{"server_runs.$.status": "skipped", "server_runs.$.finished_at": fin})
		return
	}

	_, _ = AppendMarker(runID, serverID, fmt.Sprintf("run started on %s %d step(s), workspace vantage-run-%s", serverID, len(steps), runID))

	runEnv := map[string]string{}
	allSecrets := map[string]string{}
	serverFailed := false

	// A run already in flight when the licence expires finishes its remaining
	// steps. New runs are blocked at the API, but killing a workflow midway leaves
	// a server in a half-configured state, which is worse than letting it complete.
	for i, step := range steps {
		startStep(runID, serverID, i, "running")
		stepStart := time.Now()
		var res *pb.StepResult
		attempts := 0
		maxAttempts := 1
		if step.OnFailure == "retry" {
			maxAttempts = step.MaxRetries + 1
		}

		secretVals := resolveSecrets(instanceID, step.SecretRefs)
		for k, v := range secretVals {
			allSecrets[k] = v
		}

		subst := map[string]string{}
		for k, v := range runEnv {
			subst[k] = v
		}
		for k, v := range secretVals {
			subst[k] = v
		}
		cmdEnv := map[string]string{}
		for k, v := range step.Inputs {
			cmdEnv[k] = expandVars(v, subst)
		}
		for k, v := range runEnv {
			cmdEnv[k] = v
		}
		for k, v := range secretVals {
			cmdEnv[k] = v
		}

		marker := fmt.Sprintf("===== step %d/%d: %s (%s) =====", step.Order+1, len(steps), step.Name, step.Interpreter)
		offset, _ := AppendMarker(runID, serverID, marker)
		secretsSlice := secretValues(secretVals)

		commandID := uuid.New().String()
		for attempts < maxAttempts {
			attempts++
			if attempts > 1 {
				_, _ = AppendMarker(runID, serverID, fmt.Sprintf("retry %d/%d after failure", attempts-1, maxAttempts-1))
			}

			// The log writer is opened by the pod that owns this agent's
			// stream, not here: that is where the output arrives, and the mask
			// list travels with the dispatch so it is applied before anything
			// is stored.
			res = dispatchAndWait(serverID, commandID, runID, secretsSlice, &pb.RunStepCmd{
				Interpreter:    step.Interpreter,
				Script:         step.Script,
				Env:            cmdEnv,
				TimeoutSeconds: 0,
				WorkspaceId:    runID,
			})
			if res != nil && res.ExitCode == 0 {
				break
			}
		}

		exit := 1
		outEnv := map[string]string{}
		if res != nil {
			exit = res.ExitCode
			for k, v := range res.OutputEnv {
				runEnv[k] = v
				outEnv[k] = maskSecrets(v, allSecrets)
			}
		} else {
			_, _ = AppendMarker(runID, serverID, "agent did not return a result")
		}

		status := "success"
		if exit != 0 {
			status = "failed"
		}
		finishStep(runID, serverID, i, status, attempts, exit, offset, outEnv)

		dur := time.Since(stepStart).Round(time.Millisecond)
		_, _ = AppendMarker(runID, serverID, fmt.Sprintf("step %d/%d %s exit %d, %d attempt(s), %s",
			step.Order+1, len(steps), status, exit, attempts, dur))

		if exit != 0 {
			switch step.OnFailure {
			case "continue":
				_, _ = AppendMarker(runID, serverID, "on_failure=continue proceeding to next step")
			default:
				serverFailed = true
			}
			if serverFailed {
				_, _ = AppendMarker(runID, serverID, "stopping run remaining steps skipped")
				markRemainingSkipped(runID, serverID, i+1)
				break
			}
		}
	}

	DispatchCleanupWorkspace(serverID, runID)

	fin := time.Now()
	status := "success"
	if serverFailed {
		status = "failed"
	}
	_, _ = AppendMarker(runID, serverID, fmt.Sprintf("run %s in %s workspace removed",
		status, fin.Sub(now).Round(time.Millisecond)))

	maskedRunEnv := make(map[string]string, len(runEnv))
	for k, v := range runEnv {
		maskedRunEnv[k] = maskSecrets(v, allSecrets)
	}
	setServerRun(runID, srvIdx, bson.M{
		"server_runs.$.status":      status,
		"server_runs.$.finished_at": fin,
		"server_runs.$.run_env":     maskedRunEnv,
	})
}

// dispatchAndWait subscribes to the result before dispatching, because the two
// happen on different pods and an agent that answers quickly would otherwise
// publish into a channel this process had not yet joined.
func dispatchAndWait(serverID, commandID, runID string, mask []string, cmd *pb.RunStepCmd) *pb.StepResult {
	ch, done := StepResults.Await(commandID)
	defer done()

	if err := DispatchRunStep(serverID, commandID, runID, mask, cmd); err != nil {
		return &pb.StepResult{ExitCode: 1, Stderr: "[vantage] dispatch failed: " + err.Error()}
	}
	wait := time.Duration(cmd.TimeoutSeconds)*time.Second + stepDispatchGrace
	if cmd.TimeoutSeconds == 0 {
		wait = 30*time.Minute + stepDispatchGrace
	}
	select {
	case res, ok := <-ch:
		if !ok || res == nil {
			return &pb.StepResult{ExitCode: 1, Stderr: "[vantage] lost the result channel before the agent answered"}
		}
		return res
	case <-time.After(wait):
		return &pb.StepResult{ExitCode: 124, Stderr: "[vantage] timed out waiting for agent result"}
	}
}

func expandVars(v string, lookup map[string]string) string {
	return os.Expand(v, func(name string) string {
		if name == "$" {
			return "$"
		}
		return lookup[name]
	})
}

func resolveSecrets(instanceID string, refs []string) map[string]string {
	out := map[string]string{}
	for _, ref := range refs {

		parts := strings.SplitN(ref, "/", 2)
		if len(parts) != 2 {
			continue
		}
		if v, err := RevealSecret(instanceID, parts[0], parts[1]); err == nil {
			out[parts[1]] = v
		}
	}
	return out
}

func maskSecrets(s string, secrets map[string]string) string {
	for _, v := range secrets {
		if v == "" {
			continue
		}
		s = strings.ReplaceAll(s, v, "***")
	}
	return s
}

func setServerRun(runID string, srvIdx int, set bson.M) {
	ctx, cancel := wfCtx()
	defer cancel()
	_, _ = db.Col("workflow_runs").UpdateOne(ctx,
		bson.M{"run_id": runID, "server_runs.server_id": serverIDAt(runID, srvIdx)},
		bson.M{"$set": set})
}

func serverIDAt(runID string, srvIdx int) string {
	r, err := getRunByID(runID)
	if err != nil || srvIdx >= len(r.ServerRuns) {
		return ""
	}
	return r.ServerRuns[srvIdx].ServerID
}

func startStep(runID, serverID string, order int, status string) {
	now := time.Now()
	updateStep(runID, serverID, order, bson.M{
		"server_runs.$[s].steps.$[t].status":     status,
		"server_runs.$[s].steps.$[t].started_at": now,
	})
}

func finishStep(runID, serverID string, order int, status string, attempts, exit int, logOffset int64, outEnv map[string]string) {
	now := time.Now()
	updateStep(runID, serverID, order, bson.M{
		"server_runs.$[s].steps.$[t].status":      status,
		"server_runs.$[s].steps.$[t].attempts":    attempts,
		"server_runs.$[s].steps.$[t].exit_code":   exit,
		"server_runs.$[s].steps.$[t].log_offset":  logOffset,
		"server_runs.$[s].steps.$[t].output_env":  outEnv,
		"server_runs.$[s].steps.$[t].finished_at": now,
	})
}

func secretValues(m map[string]string) []string {
	out := make([]string, 0, len(m))
	for _, v := range m {
		out = append(out, v)
	}
	return out
}

func markRemainingSkipped(runID, serverID string, fromOrder int) {
	ctx, cancel := wfCtx()
	defer cancel()
	_, _ = db.Col("workflow_runs").UpdateMany(ctx,
		bson.M{"run_id": runID},
		bson.M{"$set": bson.M{"server_runs.$[s].steps.$[t].status": "skipped"}},
		options.UpdateMany().SetArrayFilters([]interface{}{
			bson.M{"s.server_id": serverID},
			bson.M{"t.order": bson.M{"$gte": fromOrder}, "t.status": "queued"},
		}),
	)
}

func updateStep(runID, serverID string, order int, set bson.M) {
	ctx, cancel := wfCtx()
	defer cancel()
	_, _ = db.Col("workflow_runs").UpdateOne(ctx,
		bson.M{"run_id": runID},
		bson.M{"$set": set},
		options.UpdateOne().SetArrayFilters([]interface{}{
			bson.M{"s.server_id": serverID},
			bson.M{"t.order": order},
		}),
	)
}

func getRunByID(runID string) (*models.WorkflowRun, error) {
	ctx, cancel := wfCtx()
	defer cancel()
	var r models.WorkflowRun
	err := db.Col("workflow_runs").FindOne(ctx, bson.M{"run_id": runID}).Decode(&r)
	if err == mongo.ErrNoDocuments {
		return nil, fmt.Errorf("run not found")
	}
	return &r, err
}

func GetRun(instanceID, runID string) (*models.WorkflowRun, error) {
	ctx, cancel := wfCtx()
	defer cancel()
	var r models.WorkflowRun
	err := db.Col("workflow_runs").FindOne(ctx, bson.M{"run_id": runID, "instance_id": instanceID}).Decode(&r)
	if err == mongo.ErrNoDocuments {
		return nil, fmt.Errorf("run not found")
	}
	return &r, err
}

func ListRuns(instanceID, workflowID string, limit int64) ([]models.WorkflowRun, error) {
	ctx, cancel := wfCtx()
	defer cancel()
	cur, err := db.Col("workflow_runs").Find(ctx, bson.M{"instance_id": instanceID, "workflow_id": workflowID},
		options.Find().SetSort(bson.D{{Key: "started_at", Value: -1}}).SetLimit(limit))
	if err != nil {
		return nil, err
	}
	defer cur.Close(ctx)
	runs := []models.WorkflowRun{}
	if err := cur.All(ctx, &runs); err != nil {
		return nil, err
	}
	return runs, nil
}

func CancelRun(instanceID, runID string) error {
	now := time.Now()
	ctx, cancel := wfCtx()
	defer cancel()
	_, err := db.Col("workflow_runs").UpdateOne(ctx,
		bson.M{"instance_id": instanceID, "run_id": runID, "status": "running"},
		bson.M{"$set": bson.M{"status": "cancelled", "finished_at": now}})
	return err
}
