package dbreads

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/benbjohnson/litestream"
	lsfile "github.com/benbjohnson/litestream/file"
	"github.com/superfly/ltx"

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

// WorkerOptions are ServeWorker's settings. Production (RunWorker) sets
// Scratch and leaves the rest zero.
type WorkerOptions struct {
	// Scratch is where the read's private directory goes: /var/tmp, the
	// unit's own PrivateTmp, which systemd removes at stop.
	Scratch string
	// QueryTimeout is QueryTimeout when zero; Grace is WorkerExitGrace.
	QueryTimeout time.Duration
	Grace        time.Duration
	// Exit ends the process (os.Exit when nil): the 5 s backstop for a step
	// that ignores its context (spec §8.3, codex 8).
	Exit func(code int)
	// Self is the hello's view of the sandbox (SelfView when nil).
	Self func() HelloReply
	// Restore is dbdoor.RestoreLocal when nil. Tests replace it.
	Restore func(ctx context.Context, c litestream.ReplicaClient, out string, target uint64) error
	// MaxRestored is the largest restored-size cap a header may claim:
	// DefaultMaxRestored when zero, and in the unit the process's own
	// RLIMIT_FSIZE (LimitFSIZE), which a box raising the cap raises too. The
	// header's files may total at most PlanBytesFactor times its cap.
	MaxRestored int64
}

type frame struct {
	typ byte
	p   []byte
}

var errAgentGone = errors.New("dbreads: the agent closed the connection")

// ServeWorker is the worker's whole job on one connection (spec §8.3): a
// hello, or a plan, its files, a restore, ready, then run or stop. One
// goroutine reads the connection from the start; EOF or an error cancels
// the one context everything runs under, and if the work has not returned
// Grace later, the process exits.
func ServeWorker(conn io.ReadWriter, opt WorkerOptions) error {
	if opt.QueryTimeout == 0 {
		opt.QueryTimeout = QueryTimeout
	}
	if opt.Grace == 0 {
		opt.Grace = WorkerExitGrace
	}
	if opt.Exit == nil {
		opt.Exit = os.Exit
	}
	if opt.Self == nil {
		opt.Self = SelfView
	}
	if opt.Restore == nil {
		opt.Restore = dbdoor.RestoreLocal
	}
	if opt.MaxRestored <= 0 {
		opt.MaxRestored = DefaultMaxRestored
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	frames := make(chan frame)
	done := make(chan struct{})
	defer close(done)
	go func() {
		defer close(frames)
		for {
			typ, p, err := readFrame(conn)
			if err != nil {
				cancel()
				select {
				case <-done:
				case <-time.After(opt.Grace):
					opt.Exit(3)
				}
				return
			}
			select {
			case frames <- frame{typ, p}:
			case <-done:
				return
			}
		}
	}()
	next := func() (frame, error) {
		select {
		case f, ok := <-frames:
			if !ok {
				return frame{}, errAgentGone
			}
			return f, nil
		case <-ctx.Done():
			return frame{}, errAgentGone
		}
	}
	end := func(code string, err error) error {
		msg := endMsg{Code: code}
		if err != nil {
			msg.Error = err.Error()
		}
		return writeJSONFrame(conn, frameEnd, msg)
	}

	f, err := next()
	if err != nil {
		return err
	}
	switch f.typ {
	case frameHello:
		var h hello
		if json.Unmarshal(f.p, &h) != nil {
			return errProtocol
		}
		return writeJSONFrame(conn, frameHelloReply, opt.Self())
	case frameHeader:
	default:
		return errProtocol
	}
	var h planHeader
	if err := json.Unmarshal(f.p, &h); err != nil || h.V != ProtocolVersion || (h.Op != OpQuery && h.Op != OpPull) {
		return errProtocol
	}
	// Every byte this worker writes is bounded before the first is
	// written: each file by its header size (receiveFile), and all of them
	// together by PlanBytesFactor times a cap no larger than the worker's
	// own (Astra round 1, finding 1). LimitFSIZE alone bounds one file.
	if err := planWithinBounds(h, opt.MaxRestored); err != nil {
		return end(CodeCheckFailed, err)
	}
	dir, err := os.MkdirTemp(opt.Scratch, "dbread-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(dir)
	root := filepath.Join(dir, "replica")
	for _, lf := range h.Files {
		path := litestream.LTXFilePath(root, lf.Level, ltx.TXID(lf.MinTXID), ltx.TXID(lf.MaxTXID))
		if err := receiveFile(next, path, lf.Size); err != nil {
			if errors.Is(err, errAgentGone) {
				return err
			}
			return end(CodeCheckFailed, err)
		}
		// litestream's restore picks files by their names and compacts them
		// by their headers, so a file whose header is not its plan entry
		// would be reported at a transaction it does not hold (Astra round 1
		// on plan 3, finding 1). A pull's files go back unrestored, so each is
		// decoded whole as well.
		if err := verifyLTX(path, lf, h.Op == OpPull); err != nil {
			return end(CodeCheckFailed, err)
		}
	}
	if len(h.Files) == 0 || h.Files[len(h.Files)-1].MaxTXID != h.Target {
		return end(CodeCheckFailed, errors.New("dbreads: the plan does not end at its target"))
	}
	if h.Op == OpPull {
		return holdAndSend(ctx, conn, next, end, root, h)
	}
	db := filepath.Join(dir, "copy.db")
	if err := opt.Restore(ctx, lsfile.NewReplicaClient(root), db, h.Target); err != nil {
		if ctx.Err() != nil {
			return errAgentGone
		}
		return end(CodeCheckFailed, err)
	}
	if st, err := os.Stat(db); err != nil || st.Size() > h.MaxRestored {
		return end(CodeTooLarge, err)
	}
	if err := writeJSONFrame(conn, frameReady, readyMsg{TXID: h.Target}); err != nil {
		return err
	}
	f, err = next()
	if err != nil {
		return err
	}
	switch f.typ {
	case frameStop:
		return nil
	case frameRun:
	default:
		return errProtocol
	}
	var run runMsg
	if err := json.Unmarshal(f.p, &run); err != nil {
		return errProtocol
	}
	qctx, qcancel := context.WithTimeout(ctx, opt.QueryTimeout)
	defer qcancel()
	out := &frameWriter{w: conn}
	qerr := dbquery.RunContext(qctx, out, db, run.SQL, dbquery.Options{Immutable: true})
	if ferr := out.flush(); ferr != nil {
		return ferr
	}
	switch {
	case qerr == nil:
		return end("", nil)
	case errors.Is(qctx.Err(), context.DeadlineExceeded):
		return end(CodeTimeout, nil)
	case ctx.Err() != nil:
		return errAgentGone
	}
	return end(CodeSQLError, qerr)
}

// holdAndSend is a pull's whole worker job after the files are received and
// checked (spec §6: the box never restores for a pull): the files are held
// in this unit's private scratch, and after the agent's check B they go back
// in the header's order, whole, as output.
func holdAndSend(ctx context.Context, conn io.Writer, next func() (frame, error), end func(string, error) error, root string, h planHeader) error {
	// ServeWorker has checked every file whole against its plan entry, and
	// the plan's end against its target.
	paths := make([]string, len(h.Files))
	for i, lf := range h.Files {
		paths[i] = litestream.LTXFilePath(root, lf.Level, ltx.TXID(lf.MinTXID), ltx.TXID(lf.MaxTXID))
	}
	if err := writeJSONFrame(conn, frameReady, readyMsg{TXID: h.Target}); err != nil {
		return err
	}
	f, err := next()
	if err != nil {
		return err
	}
	switch f.typ {
	case frameStop:
		return nil
	case frameRun:
	default:
		return errProtocol
	}
	out := &frameWriter{w: conn}
	for _, p := range paths {
		if ctx.Err() != nil {
			return errAgentGone
		}
		in, err := os.Open(p)
		if err != nil {
			return end(CodeCheckFailed, err)
		}
		_, err = io.Copy(out, in)
		in.Close()
		if err != nil {
			return err
		}
	}
	if err := out.flush(); err != nil {
		return err
	}
	return end("", nil)
}

// verifyLTX checks one received file against its plan entry: its LTX
// header must name exactly the entry's transactions. With whole, every page
// and the checksums are decoded too (the decode is the worker's, as every
// decode of tenant bytes is).
func verifyLTX(path string, lf dbdoor.LTXFile, whole bool) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	dec := ltx.NewDecoder(f)
	if whole {
		err = dec.Verify()
	} else {
		err = dec.DecodeHeader()
	}
	if err != nil {
		return err
	}
	if h := dec.Header(); uint64(h.MinTXID) != lf.MinTXID || uint64(h.MaxTXID) != lf.MaxTXID {
		return fmt.Errorf("dbreads: %s holds transactions %s", filepath.Base(path), dbdoor.Span(uint64(h.MinTXID), uint64(h.MaxTXID)))
	}
	return nil
}

// receiveFile writes exactly size bytes of file frames to path, refusing a
// file larger than its header said (spec §8.3).
func receiveFile(next func() (frame, error), path string, size int64) error {
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return err
	}
	out, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
	if err != nil {
		return err
	}
	defer out.Close()
	var got int64
	for got < size {
		f, err := next()
		if err != nil {
			return err
		}
		if f.typ != frameFile {
			return errProtocol
		}
		if got+int64(len(f.p)) > size {
			return fmt.Errorf("dbreads: %s is larger than its header said (%d bytes)", filepath.Base(path), size)
		}
		if _, err := out.Write(f.p); err != nil {
			return err
		}
		got += int64(len(f.p))
	}
	return out.Close()
}

// frameWriter frames the query's output in outputFrameBytes pieces.
type frameWriter struct {
	w   io.Writer
	buf []byte
}

func (f *frameWriter) Write(p []byte) (int, error) {
	f.buf = append(f.buf, p...)
	for len(f.buf) >= outputFrameBytes {
		if err := writeFrame(f.w, frameOutput, f.buf[:outputFrameBytes]); err != nil {
			return 0, err
		}
		f.buf = f.buf[outputFrameBytes:]
	}
	return len(p), nil
}

func (f *frameWriter) flush() error {
	if len(f.buf) == 0 {
		return nil
	}
	err := writeFrame(f.w, frameOutput, f.buf)
	f.buf = nil
	return err
}

// SelfView is what the worker can see of the box: in the unit, the loopback
// interface alone, an empty /data and no readable /etc/carlos.
//
// Interfaces come from /proc/self/net/dev, this process's own network
// namespace. net.Interfaces cannot be used: on Linux it opens an AF_NETLINK
// socket, which the unit's RestrictAddressFamilies=AF_UNIX refuses, and an
// empty list would read as "no network" (plan review finding 1). Every
// observation that fails is reported in Errors, and a probe fails on any.
func SelfView() HelloReply {
	h := HelloReply{V: ProtocolVersion}
	if raw, err := os.ReadFile("/proc/self/net/dev"); err != nil {
		h.Errors = append(h.Errors, "interfaces: "+err.Error())
	} else {
		h.Interfaces = parseNetDev(string(raw))
	}
	switch ents, err := os.ReadDir("/data"); {
	case err == nil:
		h.DataEntries = len(ents)
	case !errors.Is(err, os.ErrNotExist):
		h.Errors = append(h.Errors, "/data: "+err.Error())
	}
	if f, err := os.Open("/etc/carlos/host.env"); err == nil {
		f.Close()
		h.EtcReadable = true
	}
	return h
}

// parseNetDev reads /proc/net/dev's interface names: two header lines, then
// one "<name>: <counters>" line per interface.
func parseNetDev(s string) []string {
	var out []string
	for i, line := range strings.Split(s, "\n") {
		if i < 2 {
			continue
		}
		name, _, ok := strings.Cut(line, ":")
		if name = strings.TrimSpace(name); ok && name != "" {
			out = append(out, name)
		}
	}
	return out
}

// planWithinBounds refuses a header whose restored-size cap is not in
// (0, limit], or whose files have a negative size or total more than
// PlanBytesFactor times that cap.
func planWithinBounds(h planHeader, limit int64) error {
	if h.MaxRestored <= 0 || h.MaxRestored > limit {
		return fmt.Errorf("dbreads: the plan's restored-size cap %d is outside (0, %d]", h.MaxRestored, limit)
	}
	var total int64
	for _, f := range h.Files {
		if f.Size < 0 || f.Size > PlanBytesFactor*h.MaxRestored-total {
			return fmt.Errorf("dbreads: the plan's files exceed %d bytes", PlanBytesFactor*h.MaxRestored)
		}
		total += f.Size
	}
	return nil
}