// Command cacheprog is a GOCACHEPROG that reads build outputs from a plain
// GOCACHE on another machine, over a command such as ssh, and keeps them in a
// scratch dir (delete it after the build):
//
//	GOCACHEPROG="cacheprog -dir /tmp/x -- ssh hat-8 cacheprog -serve GOCACHE"
//
// With -serve DIR it is the remote side: it answers "get <actionhex>" lines
// from DIR. A missing or failed remote is a cache miss, so go builds locally.
package main

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

type request struct {
	ID       int64
	Command  string
	ActionID []byte
	OutputID []byte
	BodySize int64
}

type response struct {
	ID            int64
	Err           string     `json:",omitempty"`
	KnownCommands []string   `json:",omitempty"`
	Miss          bool       `json:",omitempty"`
	OutputID      []byte     `json:",omitempty"`
	Size          int64      `json:",omitempty"`
	Time          *time.Time `json:",omitempty"`
	DiskPath      string     `json:",omitempty"`
}

func main() {
	dir := flag.String("dir", "", "local scratch dir for objects")
	serveDir := flag.String("serve", "", "serve this GOCACHE dir on stdin/stdout")
	flag.Parse()
	log.SetPrefix("cacheprog: ")
	log.SetFlags(0)
	if *serveDir != "" {
		serve(*serveDir, os.Stdin, os.Stdout)
		return
	}
	if *dir == "" {
		log.Fatal("-dir is required")
	}
	c, err := newClient(*dir, flag.Args())
	if err != nil {
		log.Fatal(err)
	}
	c.run(os.Stdin, os.Stdout)
}

type client struct {
	dir string

	mu         sync.Mutex // guards remote and the counters
	remote     *bufio.ReadWriter
	gets, hits int

	outMu sync.Mutex
	out   *bufio.Writer
	enc   *json.Encoder
}

// newClient starts the remote command, if any.
func newClient(dir string, remote []string) (*client, error) {
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, err
	}
	c := &client{dir: dir}
	if len(remote) == 0 {
		return c, nil
	}
	cmd := exec.Command(remote[0], remote[1:]...)
	cmd.Stderr = os.Stderr
	w, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	r, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	if err := cmd.Start(); err != nil {
		log.Printf("remote disabled: %v", err)
		return c, nil
	}
	c.remote = bufio.NewReadWriter(bufio.NewReaderSize(r, 1<<20), bufio.NewWriter(w))
	return c, nil
}

func (c *client) run(in io.Reader, out io.Writer) {
	c.out = bufio.NewWriter(out)
	c.enc = json.NewEncoder(c.out)
	c.send(response{KnownCommands: []string{"get", "put", "close"}})
	dec := json.NewDecoder(bufio.NewReader(in))
	var wg sync.WaitGroup
	defer wg.Wait()
	for {
		var r request
		if err := dec.Decode(&r); err != nil {
			if err != io.EOF {
				log.Fatal(err)
			}
			return
		}
		switch r.Command {
		case "get":
			wg.Add(1)
			go func() { defer wg.Done(); c.send(c.get(r)) }()
		case "put":
			var body []byte
			if r.BodySize > 0 {
				if err := dec.Decode(&body); err != nil {
					log.Fatal(err)
				}
			}
			c.send(c.put(r, body))
		case "close":
			wg.Wait()
			if c.gets > 0 && os.Getenv("CACHEPROG_STATS") != "" {
				log.Printf("remote hits %d/%d", c.hits, c.gets)
			}
			c.send(response{ID: r.ID})
			return
		default:
			c.send(response{ID: r.ID, Err: "unknown command " + r.Command})
		}
	}
}

func (c *client) send(r response) {
	c.outMu.Lock()
	defer c.outMu.Unlock()
	if err := c.enc.Encode(r); err != nil {
		log.Fatal(err)
	}
	if err := c.out.Flush(); err != nil {
		log.Fatal(err)
	}
}

func (c *client) get(r request) response {
	a := hex.EncodeToString(r.ActionID)
	if res, ok := c.local(a); ok {
		res.ID = r.ID
		return res
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	c.gets++
	if c.remote == nil {
		return response{ID: r.ID, Miss: true}
	}
	res, err := c.fetch(a)
	if err != nil {
		log.Printf("remote disabled: %v", err)
		c.remote = nil
	}
	if res == nil {
		return response{ID: r.ID, Miss: true}
	}
	c.hits++
	res.ID = r.ID
	return *res
}

// fetch asks the remote for action a; nil, nil is a miss. Holds c.mu.
func (c *client) fetch(a string) (*response, error) {
	fmt.Fprintf(c.remote, "get %s\n", a)
	if err := c.remote.Flush(); err != nil {
		return nil, err
	}
	line, err := c.remote.ReadString('\n')
	if err != nil {
		return nil, err
	}
	f := strings.Fields(line)
	if len(f) == 1 && f[0] == "miss" {
		return nil, nil
	}
	if len(f) != 4 || f[0] != "hit" || !isHash(f[1]) {
		return nil, fmt.Errorf("bad reply %q", line)
	}
	size, err1 := strconv.ParseInt(f[2], 10, 64)
	t, err2 := strconv.ParseInt(f[3], 10, 64)
	if err := errors.Join(err1, err2); err != nil || size < 0 {
		return nil, fmt.Errorf("bad reply %q", line)
	}
	if err := c.writeAtomic(f[1]+"-d", io.LimitReader(c.remote, size), size); err != nil {
		return nil, err
	}
	return c.index(a, f[1], size, t)
}

func (c *client) put(r request, body []byte) response {
	if int64(len(body)) != r.BodySize {
		return response{ID: r.ID, Err: "body size mismatch"}
	}
	a, o := hex.EncodeToString(r.ActionID), hex.EncodeToString(r.OutputID)
	err := c.writeAtomic(o+"-d", bytes.NewReader(body), r.BodySize)
	var res *response
	if err == nil {
		res, err = c.index(a, o, r.BodySize, time.Now().UnixNano())
	}
	if err != nil {
		return response{ID: r.ID, Err: err.Error()}
	}
	return response{ID: r.ID, DiskPath: res.DiskPath}
}

// index records action a -> output o in the scratch dir.
func (c *client) index(a, o string, size, t int64) (*response, error) {
	entry := fmt.Sprintf("%s %d %d", o, size, t)
	if err := c.writeAtomic(a+"-a", strings.NewReader(entry), int64(len(entry))); err != nil {
		return nil, err
	}
	res, _ := c.local(a)
	return &res, nil
}

// local looks action a up in the scratch dir, which the concurrent go builds
// of one gomobile bind share.
func (c *client) local(a string) (response, bool) {
	b, err := os.ReadFile(filepath.Join(c.dir, a+"-a"))
	if err != nil {
		return response{}, false
	}
	var o string
	var size, t int64
	if _, err := fmt.Sscanf(string(b), "%s %d %d", &o, &size, &t); err != nil || !isHash(o) {
		return response{}, false
	}
	p := filepath.Join(c.dir, o+"-d")
	if st, err := os.Stat(p); err != nil || st.Size() != size {
		return response{}, false
	}
	oid, _ := hex.DecodeString(o)
	tm := time.Unix(0, t)
	return response{OutputID: oid, Size: size, Time: &tm, DiskPath: p}, true
}

func (c *client) writeAtomic(name string, r io.Reader, size int64) error {
	f, err := os.CreateTemp(c.dir, "tmp-")
	if err != nil {
		return err
	}
	defer os.Remove(f.Name())
	n, err := io.Copy(f, r)
	if err == nil && n != size {
		err = io.ErrUnexpectedEOF
	}
	if err2 := f.Close(); err == nil {
		err = err2
	}
	if err != nil {
		return err
	}
	return os.Rename(f.Name(), filepath.Join(c.dir, name))
}

func isHash(s string) bool {
	b, err := hex.DecodeString(s)
	return err == nil && len(b) == 32
}

// serve answers "get <actionhex>" from the GOCACHE dir with "miss" or
// "hit <outputhex> <size> <unixnano>" followed by the output's bytes.
func serve(dir string, in io.Reader, w io.Writer) {
	sc := bufio.NewScanner(in)
	out := bufio.NewWriterSize(w, 1<<20)
	for sc.Scan() {
		f := strings.Fields(sc.Text())
		if len(f) != 2 || f[0] != "get" || !isHash(f[1]) || serveOne(out, dir, f[1]) != nil {
			fmt.Fprintln(out, "miss")
		}
		if err := out.Flush(); err != nil {
			log.Fatal(err)
		}
	}
}

func serveOne(out io.Writer, dir, a string) error {
	// An action entry is "v1 <actionhex> <outputhex> <size> <unixnano>\n".
	b, err := os.ReadFile(filepath.Join(dir, a[:2], a+"-a"))
	if err != nil {
		return err
	}
	f := strings.Fields(string(b))
	if len(f) != 5 || f[0] != "v1" || f[1] != a || !isHash(f[2]) {
		return errors.New("bad entry")
	}
	size, err := strconv.ParseInt(f[3], 10, 64)
	if err != nil {
		return err
	}
	d, err := os.Open(filepath.Join(dir, f[2][:2], f[2]+"-d"))
	if err != nil {
		return err
	}
	defer d.Close()
	if st, err := d.Stat(); err != nil || st.Size() != size {
		return errors.New("bad output")
	}
	fmt.Fprintf(out, "hit %s %d %s\n", f[2], size, f[4])
	if _, err := io.CopyN(out, d, size); err != nil {
		log.Fatal(err) // the stream is out of sync; the client sees EOF and stops asking
	}
	return nil
}