package main

import (
	"bytes"
	"context"
	"errors"
	"flag"
	"fmt"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"log"
	"os"
	"path/filepath"

	"chirami/internal/ocr"
	"chirami/internal/ocr/fetch"
	"chirami/internal/ocr/onnx"
	"chirami/internal/paths"
)

// defaultModelDir is where `chirami setup -ocr` puts the models.
func defaultModelDir() string {
	if env := os.Getenv("CHIRAMI_OCR_MODELS"); env != "" {
		return env
	}
	return filepath.Join(paths.DataDir(), "ocr")
}

func main() {
	log.SetFlags(0)
	log.SetPrefix("chirami-ocr: ")

	modelsFlag := flag.String("models", defaultModelDir(), "path to directory containing OCR models")
	ortLibFlag := flag.String("ort-lib", "", "explicit path to libonnxruntime shared library")
	downloadFlag := flag.Bool("download", false, "ensure all models are downloaded to -models dir and exit")
	prepareFlag := flag.Bool("prepare", false, "load every model once (writing the ORT-format copies) and exit")
	flag.Parse()

	modelsDir := *modelsFlag

	if *downloadFlag {
		log.Printf("checking and downloading models into %s...", modelsDir)
		if err := fetch.EnsureModels(context.Background(), modelsDir); err != nil {
			log.Fatalf("download models: %v", err)
		}
		log.Printf("all models verified in %s", modelsDir)
		return
	}

	// setup verified the hashes; here only a cheap check that nothing is
	// missing, since hashing 480 MB would slow every start
	if missing := fetch.Absent(modelsDir); len(missing) > 0 {
		log.Fatalf("%d model file(s) missing in %s; run `chirami setup -ocr`", len(missing), modelsDir)
	}

	if err := onnx.InitRuntime(*ortLibFlag); err != nil {
		log.Fatalf("init onnxruntime: %v", err)
	}

	pipe, err := onnx.NewPipeline(modelsDir)
	if err != nil {
		log.Fatalf("initialize pipeline: %v", err)
	}
	defer pipe.Close()
	if *prepareFlag {
		return
	}

	// Main request loop: read 4-byte big-endian length + PNG on stdin, reply JSON on stdout
	for {
		req, err := ocr.ReadRequest(os.Stdin)
		if err != nil {
			if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
				return
			}
			log.Fatalf("read request: %v", err)
		}

		img, _, err := image.Decode(bytes.NewReader(req))
		if err != nil {
			_ = ocr.WriteResult(os.Stdout, ocr.Result{
				Engine: "onnx",
				Error:  fmt.Sprintf("decode image: %v", err),
			})
			continue
		}

		res, err := pipe.Recognize(img)
		if err != nil {
			res = ocr.Result{
				Engine: "onnx",
				Error:  err.Error(),
			}
		}

		if err := ocr.WriteResult(os.Stdout, res); err != nil {
			log.Fatalf("write result: %v", err)
		}
	}
}