package main

import (
	"database/sql"
	"fmt"
	"io"
	"log"
	"os"

	"github.com/jonasknobloch/mbpe"

	"go.jknobloc.com/x/dict"
	"go.jknobloc.com/x/research/entropy"
	"go.jknobloc.com/x/shelf"
	"go.jknobloc.com/x/tokenizer/bpe"

	_ "github.com/duckdb/duckdb-go/v2"
)

type NoPreTok struct{}

func (p *NoPreTok) PreTokenize(s string) []string {
	return []string{s}
}

func main() {
	o := shelf.Abs("results/runs-entropy/lesci_minipile_test_256_seed_default/gpt2_256_m000_minipile_seed/lesci.db")

	var db *sql.DB

	if database, err := initDatabase(o); err != nil {
		log.Fatal(err)
	} else {
		db = database
	}

	defer db.Close()

	db.SetMaxOpenConns(1)

	var file *os.File

	if f, err := os.Open(shelf.Abs("results/entropy/minipile/dict.txt")); err != nil {
		log.Fatal(err)
	} else {
		file = f
	}

	tok := tokenizerMBPE()

	var dict *dict.Dict[*entropy.Entry]

	if d, err := loadDict(file, tok); err != nil {
		log.Fatal(err)
	} else {
		dict = d
	}

	if err := entropy.Run(db, dict); err != nil {
		log.Fatal(err)
	}
}

func initDatabase(dsn string) (*sql.DB, error) {
	var db *sql.DB

	if database, err := sql.Open("duckdb", dsn); err != nil {
		return nil, err
	} else {
		db = database
	}

	if err := db.Ping(); err != nil {
		_ = db.Close()

		return nil, err
	}

	return db, nil
}

func loadDict(r io.Reader, tok *mbpe.Tokenizer) (*dict.Dict[*entropy.Entry], error) {
	return dict.LoadDict[*entropy.Entry](r, func(l string) (string, *entropy.Entry) {
		var s string
		var n int

		if _, err := fmt.Sscanf(l, "%s %d", &s, &n); err != nil {
			panic(err)
		}

		ids := tok.Tokenize(s)

		b := to[int, uint8](ids) // TODO that conversion is fucking ass

		return string(b), &entropy.Entry{
			Encoded:  s, // TODO encoded or decoded
			TokenIDs: b,
			LogProbs: make([]float32, len(ids)),
			Bounds:   make([]bool, len(ids)+1),
			N:        0,
		}
	})
}

func to[a bpe.Integer, b bpe.Integer](s []a) []b {
	r := make([]b, len(s))

	// TODO handle overflow

	for i, v := range s {
		r[i] = b(v)
	}

	return r
}

func tokenizerMBPE() *mbpe.Tokenizer {
	var tok *bpe.Tokenizer

	v := shelf.Abs("tokenizers/minipile/tokenizer_gpt2_256_m000_minipile/vocab.json")
	m := shelf.Abs("tokenizers/minipile/tokenizer_gpt2_256_m000_minipile/merges.txt")

	if t, err := bpe.NewTokenizerFromFiles(v, m, bpe.DefaultConfig()); err != nil {
		log.Fatal(err)
	} else {
		tok = t
	}

	foo := bpe.MBPE(tok)

	foo.SetPreTokenizer(&NoPreTok{})

	return foo
}