// Package ocr reads Japanese text out of screenshots through a helper
// process, chirami-ocr, started on demand: Live Text on macOS, ONNX models on
// Linux. Keeping the engines out of process keeps chirami small while idle.
//
// The protocol on the helper's stdin/stdout: each request is a 4-byte
// big-endian length followed by that many bytes of PNG; each reply is one
// line of JSON, a Result. The helper exits when stdin closes.
package ocr

import (
	"bufio"
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// Line is one line or column of recognised text, in reading order.
type Line struct {
	Text     string `json:"text"`
	Vertical bool   `json:"vertical,omitempty"`
}

// Result is what the helper read in one image.
type Result struct {
	Lines  []Line `json:"lines"`
	Engine string `json:"engine,omitempty"` // which engine read it, for logs
	Error  string `json:"error,omitempty"`
}

// Text joins the lines the way Japanese is written: without spaces.
func (r Result) Text() string {
	var b strings.Builder
	for _, l := range r.Lines {
		b.WriteString(strings.TrimSpace(l.Text))
	}
	return b.String()
}

// HelperName is the helper binary, looked up next to the chirami executable
// first (as invoked, then with symlinks resolved, since nix links both into
// one bin directory), then on $PATH.
const HelperName = "chirami-ocr"

// FindHelper returns the helper's path.
func FindHelper() (string, error) {
	var dirs []string
	if strings.ContainsRune(os.Args[0], filepath.Separator) {
		dirs = append(dirs, filepath.Dir(os.Args[0]))
	}
	if exe, err := os.Executable(); err == nil {
		dirs = append(dirs, filepath.Dir(exe))
	}
	for _, d := range dirs {
		p := filepath.Join(d, HelperName)
		if _, err := os.Stat(p); err == nil {
			return p, nil
		}
	}
	return exec.LookPath(HelperName)
}

// Client talks to one helper process, starting it on first use and
// stopping it after Idle without requests.
type Client struct {
	Path string   // helper binary
	Args []string // e.g. the model directory
	Idle time.Duration

	mu    sync.Mutex
	cmd   *exec.Cmd
	in    io.WriteCloser
	out   *bufio.Reader
	timer *time.Timer
}

// Recognize reads the text in a PNG image.
func (c *Client) Recognize(ctx context.Context, png []byte) (Result, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.cmd == nil {
		if err := c.start(); err != nil {
			return Result{}, err
		}
	}
	if c.timer != nil {
		c.timer.Stop()
	}
	type reply struct {
		r   Result
		err error
	}
	done := make(chan reply, 1)
	go func() {
		var hdr [4]byte
		binary.BigEndian.PutUint32(hdr[:], uint32(len(png)))
		if _, err := c.in.Write(append(hdr[:], png...)); err != nil {
			done <- reply{err: err}
			return
		}
		line, err := c.out.ReadBytes('\n')
		if err != nil {
			done <- reply{err: err}
			return
		}
		var r Result
		err = json.Unmarshal(line, &r)
		done <- reply{r, err}
	}()
	select {
	case <-ctx.Done():
		c.stop() // the helper is mid-reply; a fresh one is simpler than resyncing
		return Result{}, ctx.Err()
	case rep := <-done:
		if rep.err != nil {
			c.stop()
			return Result{}, fmt.Errorf("ocr helper: %w", rep.err)
		}
		c.arm()
		if rep.r.Error != "" {
			return rep.r, errors.New(rep.r.Error)
		}
		return rep.r, nil
	}
}

// Close stops the helper.
func (c *Client) Close() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.stop()
}

func (c *Client) start() error {
	cmd := exec.Command(c.Path, c.Args...)
	cmd.Stderr = os.Stderr
	in, err := cmd.StdinPipe()
	if err != nil {
		return err
	}
	out, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("ocr helper: %w", err)
	}
	c.cmd, c.in, c.out = cmd, in, bufio.NewReader(out)
	return nil
}

func (c *Client) arm() {
	if c.Idle <= 0 {
		return
	}
	cmd := c.cmd
	c.timer = time.AfterFunc(c.Idle, func() {
		c.mu.Lock()
		defer c.mu.Unlock()
		if c.cmd == cmd {
			c.stop()
		}
	})
}

// stop closes stdin, which ends the helper, and reaps it.
func (c *Client) stop() {
	if c.cmd == nil {
		return
	}
	if c.timer != nil {
		c.timer.Stop()
	}
	c.in.Close()
	done := make(chan struct{})
	go func() { c.cmd.Wait(); close(done) }()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		c.cmd.Process.Kill()
		<-done
	}
	c.cmd, c.in, c.out, c.timer = nil, nil, nil, nil
}

// ReadRequest reads one request on the helper side; io.EOF means stdin
// closed and the helper should exit.
func ReadRequest(r io.Reader) ([]byte, error) {
	var hdr [4]byte
	if _, err := io.ReadFull(r, hdr[:]); err != nil {
		return nil, err
	}
	n := binary.BigEndian.Uint32(hdr[:])
	if n > 64<<20 {
		return nil, fmt.Errorf("request of %d bytes", n)
	}
	buf := make([]byte, n)
	_, err := io.ReadFull(r, buf)
	return buf, err
}

// WriteResult writes one reply on the helper side.
func WriteResult(w io.Writer, r Result) error {
	b, err := json.Marshal(r)
	if err != nil {
		return err
	}
	_, err = w.Write(append(b, '\n'))
	return err
}