package pi

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"regexp"
	"sort"
	"strings"
	"time"

	"github.com/distribution/reference"
)

var (
	ErrConfirmationRequired = errors.New("update requires --yes after reviewing the planned Pi version")
	ErrActiveSessions       = errors.New("active sessions must be drained before updating Pi; use --drain only after they are complete")
	ErrInterruptedUpdate    = errors.New("a previous Pi update is incomplete; run pi rollback --yes before starting another update")
	ErrInvalidRequest       = errors.New("invalid Pi lifecycle request")
	versionPattern          = regexp.MustCompile(`^[0-9]+(?:\.[0-9]+){1,3}(?:[-+][0-9A-Za-z.-]+)?$`)
)

// Source chooses whether the candidate is built from this checkout or pulled from an immutable image.
type Source string

const (
	BuildSource Source = "build"
	PullSource  Source = "pull"
)

// Request contains only non-secret operator inputs.
type Request struct {
	StatePath string
	Version   string
	Source    Source
	Image     string
	Confirm   bool
	Drain     bool
}

// Result summarizes the completed, failed, or recovered transaction without command output.
type Result struct {
	Phase     Phase
	StatePath string
}

// Update performs a recoverable core-only Pi update using the default Compose command layout.
func Update(ctx context.Context, runner Runner, request Request) (result Result, retErr error) {
	lock, err := acquireLock(request.StatePath)
	if err != nil {
		return Result{StatePath: request.StatePath}, err
	}
	defer lock.Release()
	if request.StatePath == "" {
		return Result{}, errors.New("update state path is required")
	}
	if !request.Confirm {
		return Result{StatePath: request.StatePath}, ErrConfirmationRequired
	}
	if !versionPattern.MatchString(request.Version) {
		return Result{StatePath: request.StatePath}, fmt.Errorf("%w: Pi version must be an explicit pinned version", ErrInvalidRequest)
	}
	if request.Source == "" { return Result{StatePath: request.StatePath}, fmt.Errorf("%w: Pi update requires an explicit source: build or pull", ErrInvalidRequest) }
	if request.Source != BuildSource && request.Source != PullSource {
		return Result{StatePath: request.StatePath}, fmt.Errorf("%w: Pi update source must be build or pull", ErrInvalidRequest)
	}
	if request.Source == PullSource {
		canonical, err := canonicalDigestReference(request.Image)
		if err != nil { return Result{StatePath: request.StatePath}, fmt.Errorf("%w: %v", ErrInvalidRequest, err) }
		request.Image = canonical
	}
	if old, err := readState(request.StatePath); err == nil && old.Phase != PhaseVerified && old.Phase != PhaseRolledBack && old.Phase != PhaseNoop {
		return Result{StatePath: request.StatePath}, ErrInterruptedUpdate
	} else if err != nil && !errors.Is(err, os.ErrNotExist) {
		return Result{StatePath: request.StatePath}, err
	}
	if err := setMaintenance(ctx, runner, true); err != nil { return Result{StatePath: request.StatePath}, err }
	defer func() {
		if clearErr := setMaintenance(context.Background(), runner, false); clearErr != nil {
			result = Result{Phase: PhaseFailed, StatePath: request.StatePath}
			if retErr == nil { retErr = errors.New("maintenance admission gate could not be cleared: recovery required")
			} else { retErr = fmt.Errorf("%w; maintenance admission gate could not be cleared: recovery required", retErr) }
		}
	}()

	running, err := activeSessions(ctx, runner)
	if err != nil {
		return Result{StatePath: request.StatePath}, err
	}
	if running {
		if !request.Drain {
			return Result{StatePath: request.StatePath}, ErrActiveSessions
		}
		for attempts := 0; attempts < 30; attempts++ {
			running, err = activeSessions(ctx, runner)
			if err != nil {
				return Result{StatePath: request.StatePath}, err
			}
			if !running {
				break
			}
			time.Sleep(time.Second)
		}
		if running {
			return Result{StatePath: request.StatePath}, ErrActiveSessions
		}
	}
	if err := Doctor(ctx, runner); err != nil {
		return Result{StatePath: request.StatePath}, err
	}

	currentVersion, err := Status(ctx, runner)
	if err != nil {
		return Result{StatePath: request.StatePath}, err
	}
	if currentVersion == request.Version {
		return Result{Phase: PhaseNoop, StatePath: request.StatePath}, nil
	}
	configured, err := renderedCore(ctx, runner)
	if err != nil {
		return Result{StatePath: request.StatePath}, err
	}
	previous, err := runningImage(ctx, runner, configured.Reference)
	if err != nil {
		return Result{StatePath: request.StatePath}, err
	}
	previous.ConfigurationSHA = configured.ConfigurationSHA
	state := State{Phase: PhasePreflight, Target: Target{Version: request.Version, Source: sourceValue(request)}, Previous: previous}
	if err := writeState(request.StatePath, state); err != nil {
		return Result{StatePath: request.StatePath}, err
	}

	state.Phase = PhaseBuilding
	if err := writeState(request.StatePath, state); err != nil { return Result{Phase: PhaseFailed, StatePath: request.StatePath}, err }
	if err := prepareCandidate(ctx, runner, request, previous.Reference); err != nil {
		return compensate(ctx, runner, request.StatePath, state, err)
	}
	if request.Drain {
		running, err = activeSessions(ctx, runner)
		if err != nil {
			return compensate(ctx, runner, request.StatePath, state, err)
		}
		if running {
			return compensate(ctx, runner, request.StatePath, state, ErrActiveSessions)
		}
	}
	if err := recreateCore(ctx, runner); err != nil {
		return compensate(ctx, runner, request.StatePath, state, err)
	}
	state.Phase = PhaseRecreated
	state.Candidate, err = runningImage(ctx, runner, previous.Reference)
	if err != nil { return compensate(ctx, runner, request.StatePath, state, err) }
	if err := writeState(request.StatePath, state); err != nil {
		return compensate(ctx, runner, request.StatePath, state, err)
	}
	if err := verifyCandidate(ctx, runner, request.Version, previous); err != nil {
		return compensate(ctx, runner, request.StatePath, state, err)
	}
	state.Phase, state.Error = PhaseVerified, ""
	if err := writeState(request.StatePath, state); err != nil {
		return compensate(ctx, runner, request.StatePath, state, err)
	}
	return Result{Phase: PhaseVerified, StatePath: request.StatePath}, nil
}

// Rollback restores the image recorded in durable update state. It is safe for interrupted runs.
func Rollback(ctx context.Context, runner Runner, statePath string, confirm bool) (result Result, retErr error) {
	lock, err := acquireLock(statePath)
	if err != nil {
		return Result{StatePath: statePath}, err
	}
	defer lock.Release()
	if !confirm {
		return Result{StatePath: statePath}, ErrConfirmationRequired
	}
	if err := setMaintenance(ctx, runner, true); err != nil { return Result{StatePath: statePath}, err }
	defer func() {
		if clearErr := setMaintenance(context.Background(), runner, false); clearErr != nil {
			result = Result{Phase: PhaseFailed, StatePath: statePath}
			if retErr == nil { retErr = errors.New("maintenance admission gate could not be cleared: recovery required")
			} else { retErr = fmt.Errorf("%w; maintenance admission gate could not be cleared: recovery required", retErr) }
		}
	}()
	if active, err := activeSessions(ctx, runner); err != nil { return Result{StatePath: statePath}, err
	} else if active { return Result{StatePath: statePath}, ErrActiveSessions }
	state, err := readState(statePath)
	if err != nil {
		return Result{StatePath: statePath}, err
	}
	if err := restore(ctx, runner, state.Previous); err != nil {
		state.Phase, state.Error = PhaseFailed, "rollback failed"
		if writeErr := writeState(statePath, state); writeErr != nil { return Result{Phase: PhaseFailed, StatePath: statePath}, errors.New("rollback failed and recovery state could not be persisted") }
		return Result{Phase: PhaseFailed, StatePath: statePath}, err
	}
	state.Phase, state.Error = PhaseRolledBack, ""
	if err := writeState(statePath, state); err != nil {
		return Result{Phase: PhaseFailed, StatePath: statePath}, errors.New("rollback restored the core but recovery state could not be persisted")
	}
	return Result{Phase: PhaseRolledBack, StatePath: statePath}, nil
}

func compensate(ctx context.Context, runner Runner, statePath string, state State, cause error) (Result, error) {
	if restoreErr := restore(ctx, runner, state.Previous); restoreErr != nil {
		state.Phase, state.Error = PhaseFailed, "candidate verification and automatic rollback failed"
		if writeErr := writeState(statePath, state); writeErr != nil {
			return Result{Phase: PhaseFailed, StatePath: statePath}, fmt.Errorf("update failed and rollback proof failed; recovery state could not be persisted")
		}
		return Result{Phase: PhaseFailed, StatePath: statePath}, fmt.Errorf("update failed; automatic rollback also failed: recovery required")
	}
	state.Phase, state.Error = PhaseRolledBack, ""
	if writeErr := writeState(statePath, state); writeErr != nil {
		return Result{Phase: PhaseFailed, StatePath: statePath}, fmt.Errorf("previous core image was restored but recovery state write failed: recovery required")
	}
	return Result{Phase: PhaseRolledBack, StatePath: statePath}, fmt.Errorf("update failed; previous core image was restored")
}

func recordFailure(path string, state State, label string, cause error) error {
	state.Phase, state.Error = PhaseFailed, label
	if err := writeState(path, state); err != nil { return fmt.Errorf("%w; recovery state write failed", cause) }
	return cause
}

func sourceValue(request Request) string {
	if request.Source == PullSource {
		return request.Image
	}
	return string(BuildSource)
}

func canonicalDigestReference(value string) (string, error) {
	if strings.Contains(value, "://") || strings.ContainsAny(value, "?#") || strings.Contains(value, "@") && strings.Contains(strings.Split(value, "@")[0], ":") && strings.Contains(strings.Split(value, "@")[0], "//") {
		return "", errors.New("pulled Pi image must be a credential-free canonical sha256 digest reference")
	}
	parsed, err := reference.ParseAnyReference(value)
	if err != nil { return "", errors.New("pulled Pi image must be a valid canonical sha256 digest reference") }
	canonical, ok := parsed.(reference.Canonical)
	if !ok || canonical.Digest().Algorithm().String() != "sha256" || len(canonical.Digest().Encoded()) != 64 {
		return "", errors.New("pulled Pi image must use an immutable sha256 digest")
	}
	return reference.FamiliarString(canonical), nil
}

func setMaintenance(ctx context.Context, runner Runner, enabled bool) error {
	path := "deactivate"
	if enabled { path = "activate" }
	args := append([]string{"exec", "-T", "core", "curl", "-fsS", "-X", "POST"}, internalIdentityHeaders...)
	args = append(args, "http://127.0.0.1:8787/internal/maintenance/"+path)
	result, err := runCompose(ctx, runner, args...)
	if err != nil { return commandError("maintenance admission gate", result, err) }
	var status struct { Active bool `json:"active"`; Admissions int `json:"admissions"` }
	if json.Unmarshal([]byte(result.Stdout), &status) != nil || status.Active != enabled || status.Admissions != 0 { return errors.New("maintenance admission gate did not acknowledge a quiescent state") }
	return nil
}

func activeSessions(ctx context.Context, runner Runner) (bool, error) {
	args := append([]string{"exec", "-T", "core", "curl", "-fsS"}, internalIdentityHeaders...)
	args = append(args, "http://127.0.0.1:8787/sessions?scope=all")
	result, err := runCompose(ctx, runner, args...)
	if err != nil {
		return false, commandError("active-session check", result, err)
	}
	var payload struct { Sessions []struct {
		Status   string `json:"status"`
		Archived bool   `json:"archived"`
	} `json:"sessions"` }
	if err := json.Unmarshal([]byte(result.Stdout), &payload); err != nil {
		return false, errors.New("active-session check returned invalid session data")
	}
	for _, session := range payload.Sessions {
		if !session.Archived && session.Status != "finalized" && session.Status != "closed" {
			return true, nil
		}
	}
	return false, nil
}

func runningImage(ctx context.Context, runner Runner, reference string) (Image, error) {
	container, err := runCompose(ctx, runner, "ps", "-q", "core")
	if err != nil || strings.TrimSpace(container.Stdout) == "" {
		return Image{}, commandError("running core image check", container, err)
	}
	id := strings.TrimSpace(container.Stdout)
	image, err := runner.Run(ctx, []string{"inspect", "--format", "{{.Image}}", id}, nil)
	if err != nil || strings.TrimSpace(image.Stdout) == "" {
		return Image{}, commandError("running core image check", image, err)
	}
	mounts, err := runner.Run(ctx, []string{"inspect", "--format", "{{json .Mounts}}", id}, nil)
	if err != nil {
		return Image{}, commandError("core volume check", mounts, err)
	}
	var raw []struct {
		Type string `json:"Type"`; Name string `json:"Name"`; Source string `json:"Source"`
		Destination string `json:"Destination"`; RW bool `json:"RW"`; Mode string `json:"Mode"`
		Propagation string `json:"Propagation"`; Driver string `json:"Driver"`
	}
	if err := json.Unmarshal([]byte(mounts.Stdout), &raw); err != nil {
		return Image{}, errors.New("core returned invalid persistence mount data")
	}
	if len(raw) == 0 {
		return Image{}, errors.New("core has no persistence mounts to preserve")
	}
	contract := make([]Mount, 0, len(raw))
	volumes := make([]string, 0, len(raw))
	for _, mount := range raw {
		if mount.Type == "" || mount.Source == "" || mount.Destination == "" { return Image{}, errors.New("core returned incomplete persistence mount data") }
		contract = append(contract, Mount{Type: mount.Type, Name: mount.Name, SourceSHA256: mountSourceHash(mount.Source), Destination: mount.Destination, RW: mount.RW, Options: strings.Join([]string{mount.Mode, mount.Propagation, mount.Driver}, "\x00")})
		if mount.Type == "volume" && mount.Name != "" {
			volumes = append(volumes, mount.Name)
		}
	}
	return Image{ID: strings.TrimSpace(image.Stdout), Reference: reference, Volumes: volumes, Mounts: contract, MountFingerprint: mountFingerprint(contract)}, nil
}

func prepareCandidate(ctx context.Context, runner Runner, request Request, reference string) error {
	if request.Source == BuildSource {
		result, err := runCompose(ctx, runner, "build", "--pull", "--build-arg", "PI_VERSION="+request.Version, "core")
		if err != nil {
			return commandError("Pi image build", result, err)
		}
		return nil
	}
	pull, err := runner.Run(ctx, []string{"pull", request.Image}, nil)
	if err != nil {
		return commandError("Pi image pull", pull, err)
	}
	tag, err := runner.Run(ctx, []string{"image", "tag", request.Image, reference}, nil)
	if err != nil {
		return commandError("Pi image tag", tag, err)
	}
	return nil
}

func recreateCore(ctx context.Context, runner Runner) error {
	result, err := runCompose(ctx, runner, "up", "--detach", "--wait", "--wait-timeout", "45", "--no-deps", "--force-recreate", "core")
	if err != nil {
		return commandError("core recreation", result, err)
	}
	return nil
}

func verifyCandidate(ctx context.Context, runner Runner, wanted string, previous Image) error {
	health, err := runCompose(ctx, runner, "exec", "-T", "core", "curl", "-fsS", "http://127.0.0.1:8787/health")
	if err != nil {
		return commandError("core health check", health, err)
	}
	version, err := Status(ctx, runner)
	if err != nil {
		return err
	}
	if version != wanted {
		return errors.New("candidate Pi version does not match requested pinned version")
	}
	if err := Test(ctx, runner); err != nil {
		return err
	}
	configured, err := renderedCore(ctx, runner)
	if err != nil {
		return err
	}
	if configured.ConfigurationSHA != previous.ConfigurationSHA {
		return errors.New("external endpoint configuration changed during Pi update")
	}
	after, err := runningImage(ctx, runner, configured.Reference)
	if err != nil {
		return err
	}
	if !sameMounts(previous.Mounts, after.Mounts) {
		return errors.New("core persistence mount contract changed during Pi update")
	}
	return nil
}

func restore(ctx context.Context, runner Runner, previous Image) error {
	tag, err := runner.Run(ctx, []string{"image", "tag", previous.ID, previous.Reference}, nil)
	if err != nil {
		return commandError("rollback image restore", tag, err)
	}
	if err := recreateCore(ctx, runner); err != nil {
		return err
	}
	configured, err := renderedCore(ctx, runner)
	if err != nil {
		return err
	}
	after, err := runningImage(ctx, runner, configured.Reference)
	if err != nil {
		return err
	}
	if after.ID != previous.ID {
		return errors.New("rollback core image does not match recorded previous image")
	}
	if configured.ConfigurationSHA != previous.ConfigurationSHA {
		return errors.New("external endpoint configuration drift prevents rollback proof")
	}
	if !sameMounts(previous.Mounts, after.Mounts) {
		return errors.New("core persistence mount contract changed during rollback")
	}
	if err := Doctor(ctx, runner); err != nil {
		return err
	}
	if err := Test(ctx, runner); err != nil {
		return err
	}
	return nil
}

func nonEmptyLines(text string) []string {
	var values []string
	for _, value := range strings.Split(text, "\n") {
		if value = strings.TrimSpace(value); value != "" {
			values = append(values, value)
		}
	}
	sort.Strings(values)
	return values
}
func sameStrings(left, right []string) bool {
	left, right = append([]string(nil), left...), append([]string(nil), right...)
	sort.Strings(left)
	sort.Strings(right)
	return strings.Join(left, "\x00") == strings.Join(right, "\x00")
}

func sameMounts(left, right []Mount) bool {
	if len(left) != len(right) {
		return false
	}
	key := func(m Mount) string {
		return m.Type + "\x00" + m.Name + "\x00" + m.SourceSHA256 + "\x00" + m.Destination + "\x00" + fmt.Sprint(m.RW) + "\x00" + m.Options
	}
	a, b := make([]string, len(left)), make([]string, len(right))
	for i := range left {
		a[i] = key(left[i])
	}
	for i := range right {
		b[i] = key(right[i])
	}
	return sameStrings(a, b)
}