package registry

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"errors"
)

type ColdRows struct {
	Route  Route
	Timers []Timer
}

// canonicalJSON decodes stored bytes into a T and encodes them again with
// this binary. Bytes written by an older binary lack any field added since,
// so they decode with those fields empty and re-encode in today's shape: an
// unchanged row then compares equal to a fresh encoding, while a row whose
// newer fields have since been set still differs. This is the only
// normalisation applied before the ownership comparisons below. An
// undecodable body returns an error, and the caller keeps its fail-closed
// ErrClaimOwnership.
func canonicalJSON[T any](stored []byte) ([]byte, error) {
	var v T
	if err := json.Unmarshal(stored, &v); err != nil {
		return nil, err
	}
	return json.Marshal(v)
}

// withoutBuildStateCopy is rt with its build_state column cleared, for the
// claim handovers' row-identity checks (RestoreClaimRows' replay,
// RebindClaimRows). The column is the agent's copy of the bucket's
// build-state object, rewritten at every reconcile refresh even while the
// instance is stopped (its read_at moves; snapshotChurnCols). A refresh
// between a handover's read and its transaction, or before its replay, is
// not a change to the installation, and refusing on it would leave a
// handover that lost its reply unable ever to replay. Neither check writes
// the column from the caller's copy on those paths.
func withoutBuildStateCopy(rt Route) Route {
	rt.BuildState = ""
	return rt
}

// coldRowsWithoutBuildStateCopy is c with its route's build_state cleared
// (withoutBuildStateCopy).
func coldRowsWithoutBuildStateCopy(c ColdRows) ColdRows {
	c.Route = withoutBuildStateCopy(c.Route)
	return c
}

func (r *Registry) RestoreClaimRows(ctx context.Context, id, retirement string, old ColdRows, next ColdRows) error {
	if r.remote != nil {
		_, err := r.remote.call(ctx, remoteRequest{Method: "RestoreClaimRows", Id: id, Retirement: retirement, Old: old, NextValue: next}, nil)
		return err
	}

	if id == "" || retirement == "" || !old.Route.ClaimOwnership.Valid() || !next.Route.ClaimOwnership.Valid() || old.Route.Host != next.Route.Host || old.Route.ClaimIncarnation != next.Route.ClaimIncarnation || next.Route.ClaimStamp <= old.Route.ClaimStamp || len(old.Timers) != len(next.Timers) {
		return ErrClaimOwnership
	}
	before, after := old.Route, next.Route
	before.ClaimOwnership = ClaimOwnership{}
	after.ClaimOwnership = ClaimOwnership{}
	if before != after {
		return ErrClaimOwnership
	}
	for i, t := range next.Timers {
		a, b := old.Timers[i], t
		a.ClaimOwnership = ClaimOwnership{}
		b.ClaimOwnership = ClaimOwnership{}
		if a != b || !t.ClaimOwnership.Valid() || t.Host != next.Route.Host || t.ClaimIncarnation != next.Route.ClaimIncarnation || t.ClaimStamp != next.Route.ClaimStamp || t.ClaimOwner != next.Route.ClaimOwner {
			return ErrClaimOwnership
		}
	}
	if _, err := r.sql.ExecContext(ctx, `CREATE TABLE IF NOT EXISTS cold_handovers(id TEXT PRIMARY KEY,body BLOB NOT NULL)`); err != nil {
		return err
	}
	tx, err := r.beginClaimTransaction(ctx, true)
	if err != nil {
		return err
	}
	defer tx.Rollback()
	encoded, err := json.Marshal(next)
	if err != nil {
		return err
	}
	var saved []byte
	err = tx.QueryRowContext(ctx, `SELECT body FROM cold_handovers WHERE id=?`, id).Scan(&saved)
	if err == nil {
		// A body saved by a binary without today's Route fields is
		// re-encoded first, or a handover started before an upgrade could
		// never be replayed after it.
		canonical, e := canonicalJSON[ColdRows](saved)
		if e != nil || !bytes.Equal(canonical, encoded) {
			return ErrClaimOwnership
		}
		current, e := scanRoute(tx.QueryRowContext(ctx, `SELECT `+routeCols+` FROM routes WHERE host=?`, next.Route.Host).Scan)
		if e != nil || withoutBuildStateCopy(current) != withoutBuildStateCopy(next.Route) {
			return ErrClaimOwnership
		}
		for _, want := range next.Timers {
			got, e := scanTimer(tx.QueryRowContext(ctx, `SELECT `+timerCols+` FROM timers WHERE host=? AND name=?`, want.Host, want.Name))
			if e != nil || got != want {
				return ErrClaimOwnership
			}
		}
		return nil
	}
	if !errors.Is(err, sql.ErrNoRows) {
		return err
	}
	verify := func(kind, host, name, rowID string, want any) error {
		var raw []byte
		if err := tx.QueryRowContext(ctx, `SELECT archive FROM claim_row_retirements WHERE journal=? AND kind=? AND host=? AND name=? AND row_id=?`, retirement, kind, host, name, rowID).Scan(&raw); err != nil {
			return err
		}
		var archive ClaimRowArchive
		if json.Unmarshal(raw, &archive) != nil {
			return ErrClaimOwnership
		}
		contents, err := json.Marshal(want)
		if err != nil {
			return err
		}
		stored := []byte(archive.Contents)
		if kind == "route" {
			// An archive retired by a binary without today's Route
			// fields is re-encoded before the comparison (see canonicalJSON).
			if stored, err = canonicalJSON[Route](stored); err != nil {
				return ErrClaimOwnership
			}
		}
		if !bytes.Equal(contents, stored) {
			return ErrClaimOwnership
		}
		return nil
	}
	if err = verify("route", old.Route.Host, "", old.Route.ClaimRowID, old.Route); err != nil {
		return err
	}
	for _, t := range old.Timers {
		if err = verify("timer", t.Host, t.Name, t.ClaimRowID, struct {
			Timer        Timer
			Registration string
		}{t, t.Registration}); err != nil {
			return err
		}
	}
	var count int
	if err = tx.QueryRowContext(ctx, `SELECT (SELECT COUNT(*) FROM routes WHERE host=?)+(SELECT COUNT(*) FROM timers WHERE host=?)`, next.Route.Host, next.Route.Host).Scan(&count); err != nil {
		return err
	}
	if count != 0 {
		return ErrClaimOwnership
	}
	if _, err = tx.ExecContext(ctx, `INSERT INTO routes (`+routeCols+`) VALUES (`+placeholders(routeColCount)+`)`, routeArgs(next.Route)...); err != nil {
		return err
	}
	for _, t := range next.Timers {
		if _, err = tx.ExecContext(ctx, `INSERT INTO timers (`+timerCols+`) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`, t.Host, t.Name, t.Kind, t.Spec, t.Path, t.NextAt, t.LastAt, t.LastStatus, t.LastMs, t.Attempts, t.CreatedAt, t.DueAt, t.Registration, t.ClaimIncarnation, t.ClaimStamp, t.ClaimOwner, t.ClaimRowID); err != nil {
			return err
		}
	}
	if _, err = tx.ExecContext(ctx, `INSERT INTO cold_handovers(id,body) VALUES (?,?)`, id, encoded); err != nil {
		return err
	}
	return tx.Commit()
}

type coldRowsWire struct {
	Route  Route
	Timers []coldTimerWire
}
type coldTimerWire struct {
	Timer        Timer
	Registration string
}

func (r ColdRows) MarshalJSON() ([]byte, error) {
	wire := coldRowsWire{Route: r.Route}
	for _, timer := range r.Timers {
		wire.Timers = append(wire.Timers, coldTimerWire{timer, timer.Registration})
	}
	return json.Marshal(wire)
}
func (r *ColdRows) UnmarshalJSON(body []byte) error {
	var wire coldRowsWire
	if err := json.Unmarshal(body, &wire); err != nil {
		return err
	}
	r.Route = wire.Route
	r.Timers = nil
	for _, timer := range wire.Timers {
		timer.Timer.Registration = timer.Registration
		r.Timers = append(r.Timers, timer.Timer)
	}
	return nil
}