package dbreads

import (
	"context"
	"crypto/ecdh"
	"encoding/json"
	"errors"
	"io"
	"net"
	"os"
	"sync/atomic"
	"time"

	"github.com/superfly/ltx"

	"github.com/carlosframework/platform/internal/activator"
	"github.com/carlosframework/platform/internal/dbdoor"
)

// serve runs one read this agent owns, inside the owner's deadline, and
// writes its result (spec §5.5 step 5). The read slot is given back here;
// the reservation stays until the worker unit is gone (J5).
func (a *Agent) serve(parent context.Context, req Request, p Picked) {
	defer a.reads.Done()
	defer a.adm.finish(req.ID)
	// The deadline is the one picked published (spec §5.7), measured on
	// this agent's clock, not a fresh OwnerDeadline from now: an election
	// whose answer came back late has already spent part of it, and the
	// console stops waiting a minute after the published one (Astra round
	// 1, finding 3). In production a.now is time.Now, so this is
	// context.WithDeadline(parent, p.Deadline).
	started := a.now()
	left := p.Deadline.Sub(started)
	ctx, cancel := context.WithTimeout(parent, left)
	defer cancel()
	var cancelled atomic.Bool
	var res Result
	if left <= 0 {
		// Past the deadline before any work: start nothing.
		a.logf("dbreads: %s: picked, but its deadline %s had passed when the election answered", req.ID, p.Deadline.Format(time.RFC3339Nano))
		res = Result{State: StateFailed, Code: CodeInterrupted}
		if pub, err := ParseReplyKey(req.ReplyKey); err == nil {
			res = a.ended(pub, req.ID, StateFailed, CodeInterrupted, Summary{}, 0)
		}
	} else {
		go a.watchCancel(ctx, req.ID, cancel, &cancelled)
		res = a.read(ctx, req)
	}
	if cancelled.Load() {
		res.State, res.Code = StateCancelled, ""
	}
	res.Box, res.Started, res.Ended = a.cfg.Box, started, a.now()
	body, err := json.Marshal(res)
	if err == nil {
		wctx, wcancel := context.WithTimeout(context.WithoutCancel(parent), 30*time.Second)
		err = putCreateOnly(wctx, a.cfg.Store, ResultKey(req.ID), body)
		wcancel()
	}
	if err != nil {
		a.logf("dbreads: %s: write the result: %v", req.ID, err)
	}
}

// watchCancel reads cancel every CancelEvery and ends the read when it
// appears: the console's best-effort word that nobody is waiting.
func (a *Agent) watchCancel(ctx context.Context, id string, cancel context.CancelFunc, seen *atomic.Bool) {
	t := time.NewTicker(a.cfg.CancelEvery)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			if _, _, err := a.cfg.Store.Get(ctx, CancelKey(id)); err == nil {
				seen.Store(true)
				cancel()
				return
			}
		}
	}
}

// ended seals the summary into a result.
func (a *Agent) ended(pub *ecdh.PublicKey, id, state, code string, s Summary, last int64) Result {
	sealed, err := SealSummary(pub, id, s)
	if err != nil {
		a.logf("dbreads: %s: seal the summary: %v", id, err)
	}
	return Result{State: state, Code: code, LastSeq: last, Sealed: sealed}
}

// read is today's "changed while it was being read" retry, box side: one
// fresh plan, with both checks again, then changed-twice.
func (a *Agent) read(ctx context.Context, req Request) Result {
	pub, err := ParseReplyKey(req.ReplyKey)
	if err != nil {
		return Result{State: StateFailed, Code: CodeInterrupted}
	}
	at, _ := req.AtTime()
	for attempt := 0; attempt < 2; attempt++ {
		if attempt > 0 {
			// The first attempt's worker may still be stopping: keep its
			// reservation, and wait for its unit before the next one (J5).
			a.adm.retry(req.ID, attempt)
			if err := a.adm.settle(ctx, settlePoll); err != nil {
				return a.ended(pub, req.ID, StateFailed, CodeInterrupted, Summary{}, 0)
			}
		}
		if res, retry := a.attempt(ctx, req, pub, at); !retry {
			return res
		}
	}
	// "Changed twice" is itself the replica's word: it waits for check B
	// like every other replica-derived ending (Astra round 1, finding 2).
	res, _ := a.afterCheckB(ctx, req, pub, CodeChangedTwice, Summary{})
	return res
}

// afterCheckB is how every ending that carries anything derived from the
// replica leaves a read: a failed restore's error text, the restored size,
// the oldest point kept, "nothing replicated", "no room", "changed twice".
// Check B runs first (spec §8.2), and when it fails the details are
// discarded: the read ends refused with the identity code alone, so a
// requester never learns anything about another instance's replica read
// under a reused prefix (Astra round 1, finding 2).
func (a *Agent) afterCheckB(ctx context.Context, req Request, pub *ecdh.PublicKey, code string, s Summary) (Result, bool) {
	_, bcode, err := a.check(ctx, req)
	switch {
	case err != nil:
		a.logf("dbreads: %s: check B before a %s ending: %v", req.ID, code, err)
		return a.ended(pub, req.ID, StateFailed, CodeInterrupted, Summary{}, 0), false
	case bcode != "":
		return a.ended(pub, req.ID, StateRefused, bcode, Summary{}, 0), false
	}
	return a.ended(pub, req.ID, StateFailed, code, s, 0), false
}

// settlePoll is how often a retry rereads the unit listing while it waits.
var settlePoll = 200 * time.Millisecond

// attempt is the identity bracket (spec §8.2): check A, plan, stream every
// planned file into a worker that restores and checks it, check B after the
// last byte has left the bucket, and only then run. No output exists before
// check B passes. retry is true when a planned file vanished or changed.
func (a *Agent) attempt(ctx context.Context, req Request, pub *ecdh.PublicKey, at time.Time) (Result, bool) {
	fail := func(code string, s Summary) (Result, bool) {
		return a.ended(pub, req.ID, StateFailed, code, s, 0), false
	}
	// failB is fail for an ending derived from the replica: check B first.
	failB := func(code string, s Summary) (Result, bool) { return a.afterCheckB(ctx, req, pub, code, s) }
	refuse := func(code string) (Result, bool) {
		return a.ended(pub, req.ID, StateRefused, code, Summary{}, 0), false
	}
	attemptStart := time.Now()
	pctx, pcancel := context.WithTimeout(ctx, a.cfg.PlanBudget)
	defer pcancel()

	rec, code, err := a.check(pctx, req) // check A
	if err != nil {
		a.logf("dbreads: %s: check A: %v", req.ID, err)
		return fail(CodeInterrupted, Summary{})
	}
	if code != "" {
		return refuse(code)
	}
	rd, err := a.cfg.OpenReplica(pctx, activator.HibernatePrefix(a.cfg.HibernatePrefix, rec.Host))
	if err != nil {
		a.logf("dbreads: %s: open the replica: %v", req.ID, err)
		return fail(CodeInterrupted, Summary{})
	}
	plan, err := dbdoor.PlanFrom(pctx, rd, at)
	var before *dbdoor.BeforeHistoryError
	switch {
	case errors.Is(err, dbdoor.ErrNothingReplicated):
		return failB(CodeNothingReplicated, Summary{})
	case errors.As(err, &before):
		return failB(CodeHistoryGone, Summary{Oldest: before.Oldest})
	case errors.Is(err, dbdoor.ErrHistoryGone):
		return failB(CodeHistoryGone, Summary{})
	case errors.Is(err, os.ErrNotExist):
		return Result{}, true
	case err != nil:
		a.logf("dbreads: %s: plan: %v", req.ID, err)
		return fail(CodeInterrupted, Summary{})
	}
	// A pull's cap is on the files it ships, a query's on the copy it
	// restores (spec §5.7). Both endings are the replica's word, so both wait
	// for check B (Astra round 1, finding 2).
	pull := req.Op == OpPull
	switch {
	case pull && plan.LTXBytes > a.cfg.MaxPull:
		return failB(CodeTooLarge, Summary{Restored: plan.LTXBytes, Cap: a.cfg.MaxPull})
	case !pull && plan.RestoredBytes > a.cfg.MaxRestored:
		return failB(CodeTooLarge, Summary{Restored: plan.RestoredBytes, Cap: a.cfg.MaxRestored})
	}
	need := plan.LTXBytes + plan.RestoredBytes
	if pull {
		need = plan.LTXBytes // the worker holds the files and restores nothing
	}
	if err := a.adm.grow(req.ID, need); err != nil {
		return failB(CodeNoRoom, Summary{})
	}
	conn, err := a.dial(pctx)
	if err != nil {
		a.logf("dbreads: %s: the worker did not start: %v", req.ID, err)
		a.askProbe()
		return fail(CodeInterrupted, Summary{})
	}
	defer conn.Close()
	defer context.AfterFunc(ctx, func() { conn.Close() })()
	stopPlanBound := context.AfterFunc(pctx, func() { conn.Close() })
	defer stopPlanBound()

	if err := writeJSONFrame(conn, frameHeader, planHeader{V: ProtocolVersion, Op: req.Op, Files: plan.Planned,
		Target: plan.TXID, MaxRestored: a.cfg.MaxRestored}); err != nil {
		closeThenProbe(a, conn)
		return fail(CodeInterrupted, Summary{})
	}
	for _, f := range plan.Planned {
		vanished, err := sendFile(pctx, conn, rd, f)
		if vanished {
			return Result{}, true
		}
		if err != nil {
			// The worker went before ready: its unit may have failed to
			// start after the socket accepted (spec §7.1; plan review
			// finding 15). Probe once this connection is closed.
			closeThenProbe(a, conn)
			return fail(CodeInterrupted, Summary{})
		}
	}
	typ, p, err := readFrame(conn)
	switch {
	case err != nil:
		closeThenProbe(a, conn)
		return fail(CodeInterrupted, Summary{})
	case typ == frameEnd:
		var e endMsg
		json.Unmarshal(p, &e)
		if e.Code != CodeTooLarge {
			e.Code = CodeCheckFailed
		}
		return failB(e.Code, Summary{Error: e.Error, Restored: plan.RestoredBytes, Cap: a.cfg.MaxRestored})
	}
	var ready readyMsg
	if typ != frameReady || json.Unmarshal(p, &ready) != nil || ready.TXID != plan.TXID {
		return failB(CodeCheckFailed, Summary{})
	}
	if _, code, err := a.check(pctx, req); err != nil || code != "" { // check B
		writeFrame(conn, frameStop, nil)
		if err != nil {
			return fail(CodeInterrupted, Summary{})
		}
		return refuse(code)
	}
	stopPlanBound() // the run is bounded by the worker and the owner's deadline
	stream, size, max := StreamRows, ChunkBytes, a.cfg.MaxOutput
	if pull {
		stream, size, max = StreamLTX, PullChunkBytes, plan.LTXBytes
	}
	// Phase timings, for operators and for plan 3's throughput measurement
	// (plan review finding 26): a query's first phase plans and restores, a
	// pull's plans and has the worker verify and hold the files; then the
	// output's sealing and confirmed creates.
	held := "restored to"
	if pull {
		held = "held through"
	}
	a.logf("dbreads: %s: planned %d files (%d bytes) and %s %d in %s", req.ID, plan.Files, plan.LTXBytes, held, plan.TXID, time.Since(attemptStart).Round(time.Millisecond))
	runStart := time.Now()
	if err := writeJSONFrame(conn, frameRun, runMsg{SQL: req.SQL}); err != nil {
		return fail(CodeInterrupted, Summary{})
	}
	cw := newChunkWriter(ctx, a.cfg.Store, req.ID, pub, stream, a.cfg.ChunkEvery, max, size)
	endCode, sqlErr, rerr := relay(conn, cw)
	last, n, sum, werr := cw.close()
	a.logf("dbreads: %s: output %d bytes in %d confirmed chunks in %s", req.ID, n, last, time.Since(runStart).Round(time.Millisecond))
	s := Summary{TXID: plan.TXID, At: plan.At, Bytes: n, SHA256: sum, Error: sqlErr}
	if pull {
		s.Files = plan.Planned
		// A worker that sends more or less than it held is a check
		// failure, never an output cap.
		if werr == nil && rerr == nil && endCode == "" && (cw.capped() || n != plan.LTXBytes) {
			return a.ended(pub, req.ID, StateFailed, CodeCheckFailed, Summary{}, last), false
		}
	}
	switch {
	case werr != nil:
		a.logf("dbreads: %s: write output: %v", req.ID, werr)
		return a.ended(pub, req.ID, StateFailed, CodeInterrupted, Summary{}, last), false
	case cw.capped():
		return a.ended(pub, req.ID, StateDone, CodeOutputCap, s, last), false
	case rerr != nil:
		return a.ended(pub, req.ID, StateFailed, CodeInterrupted, Summary{}, last), false
	case endCode == "":
		return a.ended(pub, req.ID, StateDone, "", s, last), false
	case endCode == CodeSQLError, endCode == CodeTimeout:
		return a.ended(pub, req.ID, StateFailed, endCode, s, last), false
	}
	return a.ended(pub, req.ID, StateFailed, CodeInterrupted, Summary{}, last), false
}

// sendFile streams one planned file, whole, from the bucket into the
// worker. A file that is gone, or not the size the plan listed, vanished
// under the plan: the caller re-plans.
func sendFile(ctx context.Context, w io.Writer, rd dbdoor.ReplicaReader, f dbdoor.LTXFile) (bool, error) {
	rc, err := rd.OpenLTXFile(ctx, f.Level, ltx.TXID(f.MinTXID), ltx.TXID(f.MaxTXID), 0, 0)
	if errors.Is(err, os.ErrNotExist) {
		return true, nil
	}
	if err != nil {
		return false, err
	}
	defer rc.Close()
	buf := make([]byte, fileFrameBytes)
	var sent int64
	for {
		n, rerr := rc.Read(buf)
		if n > 0 {
			if sent+int64(n) > f.Size {
				return true, nil
			}
			if err := writeFrame(w, frameFile, buf[:n]); err != nil {
				return false, err
			}
			sent += int64(n)
		}
		if rerr == io.EOF {
			break
		}
		if errors.Is(rerr, os.ErrNotExist) {
			return true, nil
		}
		if rerr != nil {
			return false, rerr
		}
	}
	return sent != f.Size, nil
}

// relay moves the worker's output into the chunk writer until the worker's
// end frame, or until the cap: then the agent stops reading and closes the
// connection, which the worker sees as EOF.
func relay(conn net.Conn, cw *chunkWriter) (string, string, error) {
	for {
		typ, p, err := readFrame(conn)
		if err != nil {
			return "", "", err
		}
		switch typ {
		case frameOutput:
			if _, err := cw.Write(p); errors.Is(err, errCapped) {
				return "", "", nil
			} else if err != nil {
				return "", "", err
			}
		case frameEnd:
			var e endMsg
			if json.Unmarshal(p, &e) != nil {
				return "", "", errProtocol
			}
			return e.Code, e.Error, nil
		default:
			return "", "", errProtocol
		}
	}
}