package tldrbench

import (
	"encoding/json"
	"html/template"
	"io"
	"os"
	"sort"
	"time"
)

// ChatInfo describes the chat that was summarized, without containing any of it.
type ChatInfo struct {
	Source     string `json:"source"` // file name only
	Messages   int    `json:"messages"`
	Kept       int    `json:"kept"`
	Chunks     int    `json:"chunks"`
	InputChars int    `json:"inputChars"`
	Truncated  bool   `json:"truncated"`
}

// Run is one model writing one overview with one prompt.
type Run struct {
	Model     string  `json:"model"`
	Variant   string  `json:"variant"`
	N         int     `json:"n"`
	Seconds   float64 `json:"seconds"`
	FirstText float64 `json:"firstText"` // seconds until the first streamed text, 0 if never
	Output    string  `json:"output,omitempty"`
	Error     string  `json:"error,omitempty"`
	Score
}

// Report is everything the dashboard shows.
type Report struct {
	When     time.Time `json:"when"`
	Chat     ChatInfo  `json:"chat"`
	Threads  int       `json:"threads"`
	Facts    []string  `json:"facts,omitempty"`
	Models   []string  `json:"models"`
	Variants []string  `json:"variants"`
	Runs     []Run     `json:"runs"`
}

// WriteJSON saves the raw results so they can be re-rendered or compared later.
func (r *Report) WriteJSON(path string) error {
	raw, err := json.MarshalIndent(r, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, raw, 0o600)
}

// ReadReport loads a report written by WriteJSON.
func ReadReport(path string) (*Report, error) {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var r Report
	if err := json.Unmarshal(raw, &r); err != nil {
		return nil, err
	}
	return &r, nil
}

// Cell aggregates the runs of one model with one prompt.
type Cell struct {
	Runs, Failed              int
	Seconds, FirstText, Facts float64
	Invented                  float64
	Copied, English, CutOff   int
	HitCount                  map[string]int
	HasFacts                  bool
	Class                     string // ok, mid, no, na
	SecondsMin, SecondsMax    float64
	FactsTotal                int
	FactsPct                  int
	Flags                     []string
}

func aggregate(runs []Run) Cell {
	c := Cell{HitCount: map[string]int{}, SecondsMin: -1}
	var ok int
	for _, r := range runs {
		c.Runs++
		if r.Error != "" {
			c.Failed++
			continue
		}
		ok++
		c.Seconds += r.Seconds
		c.FirstText += r.FirstText
		c.Facts += float64(r.FactsHit)
		c.Invented += float64(len(r.Invented))
		c.FactsTotal = r.FactsTotal
		if r.Copied {
			c.Copied++
		}
		if r.EnglishIntro {
			c.English++
		}
		if r.CutOff {
			c.CutOff++
		}
		for _, h := range r.Hit {
			c.HitCount[h]++
		}
		if c.SecondsMin < 0 || r.Seconds < c.SecondsMin {
			c.SecondsMin = r.Seconds
		}
		if r.Seconds > c.SecondsMax {
			c.SecondsMax = r.Seconds
		}
	}
	if ok > 0 {
		n := float64(ok)
		c.Seconds /= n
		c.FirstText /= n
		c.Facts /= n
		c.Invented /= n
	}
	c.HasFacts = c.FactsTotal > 0
	c.Class = "na"
	if ok > 0 && c.HasFacts {
		c.FactsPct = int(100*c.Facts/float64(c.FactsTotal) + 0.5)
		switch {
		case c.FactsPct >= 80:
			c.Class = "ok"
		case c.FactsPct >= 55:
			c.Class = "mid"
		default:
			c.Class = "no"
		}
	}
	if c.Copied > 0 {
		c.Flags = append(c.Flags, "copies the chat")
	}
	if c.English > 0 {
		c.Flags = append(c.Flags, "English intro")
	}
	if c.CutOff > 0 {
		c.Flags = append(c.Flags, "may be cut off")
	}
	if c.Invented > 0 {
		c.Flags = append(c.Flags, "numbers not in the chat")
	}
	if c.Failed > 0 {
		c.Flags = append(c.Flags, "failed runs")
	}
	return c
}

type viewRow struct {
	Variant string
	Cells   []Cell
}

type viewModelFacts struct {
	Model string
	Rows  []struct {
		Variant string
		Hits    []int
		Runs    int
	}
}

type view struct {
	Report
	Rows       []viewRow
	FactGrid   []viewModelFacts
	BestModel  string
	BestVar    string
	HasText    bool
	RunsSorted []Run
}

// WriteHTML renders a self-contained dashboard: inline styles, no scripts, no
// external requests.
func (r *Report) WriteHTML(w io.Writer) error {
	v := view{Report: *r}
	cell := map[[2]string][]Run{}
	for _, run := range r.Runs {
		k := [2]string{run.Model, run.Variant}
		cell[k] = append(cell[k], run)
		if run.Output != "" {
			v.HasText = true
		}
	}
	best := -1.0
	for _, variant := range r.Variants {
		row := viewRow{Variant: variant}
		for _, model := range r.Models {
			c := aggregate(cell[[2]string{model, variant}])
			row.Cells = append(row.Cells, c)
			if c.HasFacts && c.Runs > c.Failed {
				score := c.Facts - c.Seconds/1000 // ties go to the faster one
				if score > best {
					best, v.BestModel, v.BestVar = score, model, variant
				}
			}
		}
		v.Rows = append(v.Rows, row)
	}
	if len(r.Facts) > 0 {
		for _, model := range r.Models {
			mf := viewModelFacts{Model: model}
			for _, variant := range r.Variants {
				c := aggregate(cell[[2]string{model, variant}])
				row := struct {
					Variant string
					Hits    []int
					Runs    int
				}{Variant: variant, Runs: c.Runs - c.Failed}
				for _, f := range r.Facts {
					row.Hits = append(row.Hits, c.HitCount[f])
				}
				mf.Rows = append(mf.Rows, row)
			}
			v.FactGrid = append(v.FactGrid, mf)
		}
	}
	v.RunsSorted = append(v.RunsSorted, r.Runs...)
	sort.SliceStable(v.RunsSorted, func(i, j int) bool {
		a, b := v.RunsSorted[i], v.RunsSorted[j]
		if a.Model != b.Model {
			return indexOf(r.Models, a.Model) < indexOf(r.Models, b.Model)
		}
		if a.Variant != b.Variant {
			return indexOf(r.Variants, a.Variant) < indexOf(r.Variants, b.Variant)
		}
		return a.N < b.N
	})
	return page.Execute(w, v)
}

func indexOf(list []string, s string) int {
	for i, x := range list {
		if x == s {
			return i
		}
	}
	return len(list)
}

var page = template.Must(template.New("page").Funcs(template.FuncMap{
	"inc": func(n int) int { return n + 1 },
	"f1":  func(f float64) string { return fmtFloat(f, 1) },
	"sec": func(f float64) string { return fmtFloat(f, 0) },
	"hitClass": func(hits, runs int) string {
		switch {
		case runs == 0:
			return "h0"
		case hits == runs:
			return "h2"
		case hits > 0:
			return "h1"
		}
		return "h0"
	},
}).Parse(pageHTML))