package membank

// labelharness_test.go: the K6 (glossary post-check) half of the durable labels re-measurement — the
// companion to internal/checks/labelharness_test.go. It lives here because dstFormPresent / entry are
// unexported. Env-gated exactly the same way (TM_CHECKER_LABELS=1, corpus OUT of git).
//
// For every (unit, glossary term) row of the package-6 pool it asks the REAL dstFormPresent whether an
// accepted dst form appears in the unit's final, and joins that to the human label (defect = a real miss,
// ok = the form is present). predicted_defect = NOT dstFormPresent. The package-6 baseline (metrics.json
// k6) predates the bank-quality stemmer (D39.75), so the harness measures BOTH the pre-stemmer path (inert
// stemmer == the ratified baseline) and the current path (stemmer on) — the delta is the stemmer / #10 work.
import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"textmachine/backend/internal/lang"
	"textmachine/backend/internal/text"
)

func envOrMem(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}

var memLabelsDir = envOrMem("TM_CHECKER_LABELS_DIR", "/home/ubuntu/books/gu-zhenren/labels")

type k6Pool struct {
	ID      string   `json:"id"`
	Run     string   `json:"run"`
	Chapter int      `json:"chapter"`
	Chunk   int      `json:"chunk_idx"`
	Src     string   `json:"src"`
	Dst     string   `json:"dst"`
	Forms   []string `json:"forms"`
	Status  string   `json:"status"`
	Type    string   `json:"type"`
}

type k6Label struct {
	ID    string `json:"id"`
	Label string `json:"label"`
}

type memUnit struct {
	Run     string `json:"run"`
	Chapter int    `json:"chapter"`
	Chunk   int    `json:"chunk_idx"`
	Final   string `json:"final"`
}

// stripManifestHeadingMem mirrors the checks harness: the post-check also sees final BEFORE ApplyHeading.
func stripManifestHeadingMem(final string) string {
	lines := strings.Split(final, "\n")
	if len(lines) == 0 {
		return final
	}
	rest, ok := strings.CutPrefix(strings.TrimSpace(lines[0]), "Глава ")
	if !ok {
		return final
	}
	for _, r := range rest {
		if r < '0' || r > '9' {
			return final
		}
	}
	i := 1
	for i < len(lines) && strings.TrimSpace(lines[i]) == "" {
		i++
	}
	return strings.Join(lines[i:], "\n")
}

// k6metric is the confusion count with the same defect>disputable>ok convention (here labels are per-row).
type k6metric struct{ TP, FP, FN, TN, Disp int }

func (m *k6metric) add(pred, defect bool) {
	switch {
	case pred && defect:
		m.TP++
	case pred && !defect:
		m.FP++
	case !pred && defect:
		m.FN++
	default:
		m.TN++
	}
}

// measureK6 runs dstFormPresent over every pool row against its unit's final and joins to the labels.
// When demoteSingleHan is set it models the #10 fix: a single-Han-rune key miss is demoted to
// observability (not a predicted CONFIRMED defect), exactly as production Postcheck now routes it.
func measureK6(t *testing.T, stemmer lang.TargetStemmer, demoteSingleHan bool) (k6metric, []string) {
	t.Helper()
	// corpus finals keyed by unit id.
	finals := map[string]string{}
	cb, err := os.ReadFile(filepath.Join(memLabelsDir, "raw", "corpus.jsonl"))
	if err != nil {
		t.Fatalf("read corpus: %v", err)
	}
	for _, line := range strings.Split(strings.TrimSpace(string(cb)), "\n") {
		if strings.TrimSpace(line) == "" {
			continue
		}
		var u memUnit
		if err := json.Unmarshal([]byte(line), &u); err != nil {
			t.Fatalf("parse corpus: %v", err)
		}
		finals[fmt.Sprintf("%s:%d:%d", u.Run, u.Chapter, u.Chunk)] = u.Final
	}
	// labels.
	labels := map[string]string{}
	lb, err := os.ReadFile(filepath.Join(memLabelsDir, "labels", "k6.jsonl"))
	if err != nil {
		t.Fatalf("read k6 labels: %v", err)
	}
	for _, line := range strings.Split(strings.TrimSpace(string(lb)), "\n") {
		if strings.TrimSpace(line) == "" {
			continue
		}
		var l k6Label
		if err := json.Unmarshal([]byte(line), &l); err != nil {
			t.Fatalf("parse k6 label: %v", err)
		}
		labels[l.ID] = l.Label
	}
	// pool → dstFormPresent per row, joined to labels.
	pb, err := os.ReadFile(filepath.Join(memLabelsDir, "pools", "k6_glossary.jsonl"))
	if err != nil {
		t.Fatalf("read k6 pool: %v", err)
	}
	var m k6metric
	var flaggedOK []string // fp ids (flagged as miss but label=ok) — the #10 / inflection-gap surface
	for _, line := range strings.Split(strings.TrimSpace(string(pb)), "\n") {
		if strings.TrimSpace(line) == "" {
			continue
		}
		var p k6Pool
		if err := json.Unmarshal([]byte(line), &p); err != nil {
			t.Fatalf("parse k6 pool: %v", err)
		}
		lab, ok := labels[p.ID]
		if !ok {
			continue // pool wider than the labelled sample
		}
		final := stripManifestHeadingMem(finals[fmt.Sprintf("%s:%d:%d", p.Run, p.Chapter, p.Chunk)])
		nout := text.NormalizeTargetForm(final)
		noutRunes := []rune(nout)
		outWords := lang.TokenizeWords(nout)
		var forms []string
		for _, f := range p.Forms {
			if nf := text.NormalizeTargetForm(f); nf != "" {
				forms = append(forms, nf)
			}
		}
		e := &entry{src: p.Src, dst: p.Dst, declForms: forms}
		present := dstFormPresent(e, noutRunes, outWords, stemmer)
		// The firing key for a pool row is its src (the term whose key matched). Model the #10 demote.
		predDefect := !present && !(demoteSingleHan && singleHanKeyFired(p.Src))
		switch lab {
		case "disputable":
			m.Disp++
		case "defect":
			m.add(predDefect, true)
		default:
			if predDefect {
				flaggedOK = append(flaggedOK, p.ID)
			}
			m.add(predDefect, false)
		}
	}
	return m, flaggedOK
}

func TestK6LabelsBaseline(t *testing.T) {
	if _, err := os.Stat(filepath.Join(memLabelsDir, "raw", "corpus.jsonl")); err != nil {
		if os.Getenv("TM_CHECKER_LABELS") == "1" {
			t.Fatalf("TM_CHECKER_LABELS=1 but corpus is missing: %v", err)
		}
		t.Skipf("labelled corpus absent (%v) — set TM_CHECKER_LABELS=1 to force", err)
	}

	// Pre-stemmer, pre-#10 path = the ratified package-6 baseline (metrics.json k6).
	inert := lang.TargetStemmer{}
	base, baseFP := measureK6(t, inert, false)
	t.Logf("K6 BASELINE (pre-stemmer, pre-#10 = metrics.json): tp=%d fp=%d fn=%d tn=%d disp=%d", base.TP, base.FP, base.FN, base.TN, base.Disp)
	t.Logf("K6 baseline fp ids: %v", baseFP)
	if (base != k6metric{TP: 1, FP: 14, FN: 0, TN: 243, Disp: 2}) {
		t.Errorf("k6 baseline drift: got %+v want tp1 fp14 fn0 tn243 disp2", base)
	}

	// Stemmer on (D39.75), no #10 — the intermediate landed state. Delta vs baseline = inflection-gap fps
	// the stemmer now accepts (凡人→смертных, 甲等→разряд А).
	stem := lang.NewTargetStemmer(lang.TargetChecksFor("ru"))
	mid, midFP := measureK6(t, stem, false)
	t.Logf("K6 stemmer-on, pre-#10: fp=%d (down from 14 — %v) precision=%.4f", mid.FP, diff(baseFP, midFP), prec(mid))

	// Stemmer on + #10 demote — THIS pack's after-state. The 6 single-Han-rune (转) misses demote to
	// observability, leaving the 元石/蛊虫 inflection residual (a SEED-completeness gap, not a code one).
	after, afterFP := measureK6(t, stem, true)
	t.Logf("K6 AFTER (stemmer + #10): tp=%d fp=%d fn=%d tn=%d disp=%d precision=%.4f", after.TP, after.FP, after.FN, after.TN, after.Disp, prec(after))
	t.Logf("K6 after fp ids (residual = 元石/蛊虫 inflection-gap, seed concern): %v", afterFP)
	if (after != k6metric{TP: 1, FP: 6, FN: 0, TN: 251, Disp: 2}) {
		t.Errorf("k6 after-#10 drift: got %+v want tp1 fp6 fn0 tn251 disp2", after)
	}
}

func prec(m k6metric) float64 {
	if m.TP+m.FP == 0 {
		return 0
	}
	return float64(m.TP) / float64(m.TP+m.FP)
}

// diff returns ids in a not in b (the fps removed by a change).
func diff(a, b []string) []string {
	in := map[string]bool{}
	for _, s := range b {
		in[s] = true
	}
	var out []string
	for _, s := range a {
		if !in[s] {
			out = append(out, s)
		}
	}
	return out
}