package entropy

import (
	"cmp"
	"database/sql"
	"encoding/gob"
	"os"
	"time"

	"github.com/jonasknobloch/mbpe"

	"go.jknobloc.com/x/dict"
	"go.jknobloc.com/x/shelf"
	"go.jknobloc.com/x/tui"
)

type segmentsEntry struct {
	Value []string
	OK    bool
}

const BufferSize = 1024

func Run(db *sql.DB, d *dict.Dict[*Entry]) error {
	// m := 0
	//
	// for _, v := range d.Values() {
	// 	m = max(m, len(v.TokenIDs))
	// }
	//
	// if m > BufferSize {
	// 	// TODO handle
	// }

	var rows *sql.Rows

	if r, err := db.Query(`SELECT * FROM context`); err != nil {
		return err
	} else {
		rows = r
	}

	defer rows.Close()

	bufferTokens := make([]uint8, 0, BufferSize)
	bufferLogProbs := make([]float32, 0, BufferSize)

	doc := -1

	maxTokens := 100000
	currentTokens := 0

	pb := tui.NewProgressBar("Adding LogProbs", 20, 58508191, time.Now())

	pb.Start(1*time.Second, func() int {
		return pb.Completed()
	})

	for rows.Next() {
		if currentTokens == maxTokens {
			break
		}

		currentTokens++

		var uid int
		var token int
		var logProb float32
		var pos int

		if err := rows.Scan(&uid, &token, &logProb, &pos); err != nil {
			return err
		}

		if uid != doc || len(bufferTokens) == BufferSize {
			bufferTokens = bufferTokens[:0]
			bufferLogProbs = bufferLogProbs[:0]

			doc = uid
		}

		bufferTokens = append(bufferTokens, uint8(token))
		bufferLogProbs = append(bufferLogProbs, logProb)

		n := len(bufferTokens)

		// fmt.Print(n)

		for i := range n {
			entry, ok := d.GetBytes(bufferTokens[i:n])

			if !ok {
				continue
			}

			for j := range entry.LogProbs {
				entry.LogProbs[j] += bufferLogProbs[j]
			}

			entry.N++
		}

		pb.Add(1)
	}

	if err := rows.Err(); err != nil {
		return err
	}

	// debug
	d.Sort(func(a, b dict.Entry[*Entry]) int {
		return cmp.Compare(b.Val.N, a.Val.N)
	})

	pb.Close()

	for _, e := range d.Values() {
		spikes(e)
	}

	if err := serializeSegments(d); err != nil {
		return err
	}

	return nil
}

func spikes(e *Entry) {
	for i := range len(e.LogProbs) {
		if i == 0 {
			continue
		}

		if i == 1 && e.TokenIDs[0] == 220 {
			continue
		}

		// TODO threshold

		if d := e.LogProbs[i-1] - e.LogProbs[i]; d > 0 {
			e.Bounds[i] = true
		}
	}
}

func serializeSegments(d *dict.Dict[*Entry]) error {
	r := make(map[string]segmentsEntry, d.Len())

	for _, e := range d.Values() {
		segments := make([]string, 0)

		decoded := decode(e.Encoded)

		i := 0

		for j, b := range e.Bounds {
			if b {
				segments = append(segments, encode(decoded[i:j]))

				i = j
			}
		}

		segments = append(segments, encode(decoded[i:])) // TODO or just use OK false?

		// debug
		if i == 0 {
			continue
		}

		r[e.Encoded] = segmentsEntry{
			Value: segments,
			OK:    true,
		}
	}

	var file *os.File

	if f, err := os.Create(shelf.Abs("results/entropy/minipile/segments.gob")); err != nil {
		return err
	} else {
		file = f
	}

	defer file.Close()

	enc := gob.NewEncoder(file)

	return enc.Encode(r)
}

func encode(s string) string {
	result := ""

	for _, b := range []byte(s) {
		result += mbpe.BytesChar[b]
	}

	return result
}

func decode(s string) string {
	result := make([]byte, 0)

	for _, r := range s {
		result = append(result, mbpe.CharBytes[string(r)])
	}

	return string(result)
}