package semantic

import (
	"crypto/sha256"
	_ "embed"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
)

//go:embed corpus_v1.json
var corpusV1JSON []byte

// RouteExample — one labeled utterance in the coarse-route experiment corpus.
// Every row carries provenance so no row can masquerade as independent when
// related paraphrases exist.
type RouteExample struct {
	Text       string        `json:"text"`
	Route      SemanticRoute `json:"route"`
	Source     string        `json:"source"`      // where this row came from
	SourceID   string        `json:"source_id"`   // case ID in the source corpus
	SplitGroup string        `json:"split_group"` // family/seed ID for split discipline
	Tags       []string      `json:"tags,omitempty"`
	// FastPathResolved is true when a stage-0 grammar already handles this
	// utterance. The learned router should not be measured on these unless
	// explicitly desired; they are tagged, not removed.
	FastPathResolved bool `json:"fast_path_resolved"`
	// RouterResidual marks rows that would actually reach the general routing
	// cascade after pre-route resolvers (command-prohibition, etc.) have had
	// first refusal. Static corpus evaluation cannot determine this for all
	// cases (some need dialogue state), so this is an explicit annotation
	// rather than a derived field.
	RouterResidual *bool `json:"router_residual,omitempty"`
}

// CorpusEnvelope — the versioned JSON envelope with reproducibility metadata.
type CorpusEnvelope struct {
	SchemaVersion int           `json:"schema_version"`
	Name          string        `json:"name"`
	Notes         []string      `json:"notes"`
	// Reproducibility metadata — informational, not validated against at
	// load time. The hash fields are precomputed from the source fixtures
	// at corpus-build time and recorded here so a later reader can verify
	// the corpus was built from the expected inputs.
	Reproducibility *ReproducibilityMeta `json:"reproducibility,omitempty"`
	Examples        []RouteExample       `json:"examples"`
}

// ReproducibilityMeta — dataset identity for later experiment reproduction.
type ReproducibilityMeta struct {
	// SourceFixtureHash is the SHA-256 of the concatenated source fixture
	// files used to build this corpus, hex-encoded, first 16 bytes.
	SourceFixtureHash string `json:"source_fixture_hash"`
	// ContrastGeneratorVersion identifies the transform code version.
	ContrastGeneratorVersion string `json:"contrast_generator_version"`
	// SplitAlgorithm identifies the split algorithm and version.
	SplitAlgorithm string `json:"split_algorithm"`
	// DatasetHash is the SHA-256 of the sorted example texts, hex-encoded,
	// first 16 bytes. Computed at validation time.
	DatasetHash string `json:"dataset_hash"`
	// FrozenHoldoutHash is the SHA-256 of the frozen holdout group IDs,
	// computed when the split is created.
	FrozenHoldoutHash string `json:"frozen_holdout_hash,omitempty"`
}

// SchemaVersionV1 is the version this package understands.
const SchemaVersionV1 = 1

// CorpusStats — computed from a validated corpus. Returned by ValidateCorpus
// so callers get the numbers without recomputing.
type CorpusStats struct {
	Total       int
	RouteCounts map[SemanticRoute]int
	SourceCounts map[string]int
	FastPath    int
	Residual    int
	DatasetHash string
}

// LoadCorpus returns the embedded corpus, rejecting unknown schema versions
// and failing on structural validation errors.
func LoadCorpus() ([]RouteExample, error) {
	var env CorpusEnvelope
	if err := json.Unmarshal(corpusV1JSON, &env); err != nil {
		return nil, fmt.Errorf("semantic corpus: parse: %w", err)
	}
	if env.SchemaVersion != SchemaVersionV1 {
		return nil, fmt.Errorf("semantic corpus: schema_version %d, want %d",
			env.SchemaVersion, SchemaVersionV1)
	}
	if err := ValidateCorpus(env.Examples); err != nil {
		return nil, fmt.Errorf("semantic corpus: %w", err)
	}
	return env.Examples, nil
}

// ValidateCorpus checks structural invariants: total matches, route/source
// sums, fast-path+residual, no duplicate identities, no conflicting labels.
func ValidateCorpus(exs []RouteExample) error {
	if len(exs) == 0 {
		return fmt.Errorf("corpus is empty")
	}

	routeCounts := map[SemanticRoute]int{}
	sourceCounts := map[string]int{}
		type idKey struct{ Source, SourceID, Text string }
		seen := map[idKey]bool{}
	textRoute := map[string]SemanticRoute{}

	for i, e := range exs {
		if e.Source == "" {
			return fmt.Errorf("row %d: empty source (source_id=%q)", i, e.SourceID)
		}
		if e.SourceID == "" {
			return fmt.Errorf("row %d: empty source_id (source=%q)", i, e.Source)
		}
		if e.SplitGroup == "" {
			return fmt.Errorf("row %d: empty split_group (source_id=%q)", i, e.SourceID)
		}
		if !ValidRoute(e.Route) {
			return fmt.Errorf("row %d: invalid route %q (source_id=%q)", i, e.Route, e.SourceID)
		}

		// Check for conflicting labels on identical text.
		norm := strings.TrimSpace(e.Text)

		key := idKey{Source: e.Source, SourceID: e.SourceID, Text: norm}
		if seen[key] {
			return fmt.Errorf("row %d: duplicate source+source_id+text %q:%q:%q", i, e.Source, e.SourceID, norm)
		}
		seen[key] = true

		routeCounts[e.Route]++
		sourceCounts[e.Source]++

		if prev, ok := textRoute[norm]; ok && prev != e.Route {
			return fmt.Errorf("row %d: text %q has route %q, but earlier row had %q",
				i, norm, e.Route, prev)
		}
		textRoute[norm] = e.Route
	}

	// Verify fast-path + residual == total.
	fastPath, residual := SplitCounts(exs)
	if fastPath+residual != len(exs) {
		return fmt.Errorf("fast_path(%d) + residual(%d) = %d != total(%d)",
			fastPath, residual, fastPath+residual, len(exs))
	}

	return nil
}

// CorpusStatsFrom computes the stats for a validated corpus.
func CorpusStatsFrom(exs []RouteExample) CorpusStats {
	routeCounts := map[SemanticRoute]int{}
	sourceCounts := map[string]int{}
	for _, e := range exs {
		routeCounts[e.Route]++
		sourceCounts[e.Source]++
	}
	fastPath, residual := SplitCounts(exs)

	// Deterministic dataset hash: sort texts, hash the concatenation.
	texts := make([]string, len(exs))
	for i, e := range exs {
		texts[i] = e.Text
	}
	sort.Strings(texts)
	h := sha256.Sum256([]byte(strings.Join(texts, "\n")))

	return CorpusStats{
		Total:        len(exs),
		RouteCounts:  routeCounts,
		SourceCounts: sourceCounts,
		FastPath:     fastPath,
		Residual:     residual,
		DatasetHash:  hex.EncodeToString(h[:16]),
	}
}

// ByRoute groups examples by their semantic route, for per-route inspection.
func ByRoute(exs []RouteExample) map[SemanticRoute][]RouteExample {
	m := make(map[SemanticRoute][]RouteExample)
	for _, e := range exs {
		m[e.Route] = append(m[e.Route], e)
	}
	for k := range m {
		sort.Slice(m[k], func(i, j int) bool { return m[k][i].Text < m[k][j].Text })
	}
	return m
}

// SplitCounts returns the number of fast-path-resolved vs
// general-route-required examples.
func SplitCounts(exs []RouteExample) (fastPath, residual int) {
	for _, e := range exs {
		if e.FastPathResolved {
			fastPath++
		} else {
			residual++
		}
	}
	return
}