package lib

import (
	"encoding/json"
	"fmt"
	"log"
	"sort"
	"strings"
)

// Normalize normalizes entry names within a dossier for a given category.
// Uses heuristic pre-grouping + LLM to map variant names to canonical forms.
// Updates Summary (display) and Data JSON (normalized_name, abbreviation).
// Original Type field is never modified.
// Silently returns nil if no API key is configured.
func Normalize(dossierID string, category int, progress ...func(processed, total int)) error {
	reportProgress := func(p, t int) {
		if len(progress) > 0 && progress[0] != nil {
			progress[0](p, t)
		}
	}
	if GeminiKey == "" {
		return nil
	}

	// 1. Load all entries, collect types only from entries that need normalization
	entries, err := EntryQueryOld(dossierID, category, "")
	if err != nil {
		return fmt.Errorf("load entries: %w", err)
	}

	seen := make(map[string]bool)
	var allNames []string
	for _, e := range entries {
		if e.ParentID == "" || e.Type == "lab_order" || e.Type == "" {
			continue
		}
		// FIXED(review-2026-02-28): Skip only if FULLY normalized (has both SearchKey2 AND LOINC)
		// Previously skipped on SearchKey2 alone, causing LOINC to never be populated
		if e.SearchKey2 != "" {
			var data map[string]interface{}
			json.Unmarshal([]byte(e.Data), &data)
			if loinc, ok := data["loinc"].(string); ok && loinc != "" {
				continue // fully normalized
			}
			// Has SearchKey2 but no LOINC - needs normalization
		}
		if !seen[e.Type] {
			seen[e.Type] = true
			allNames = append(allNames, e.Type)
		}
	}

	if len(allNames) == 0 {
		log.Printf("normalize: all entries already normalized")
		return nil
	}

	// 2. Pre-group by heuristic key (strip POCT, specimen suffixes, normalize case)
	groups := make(map[string][]string) // cleanKey → [original names]
	for _, name := range allNames {
		key := normalizeKey(name)
		groups[key] = append(groups[key], name)
	}

	// Send just the group keys to LLM
	keys := make([]string, 0, len(groups))
	for k := range groups {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	log.Printf("normalize: %d unique types → %d groups after pre-grouping", len(allNames), len(keys))

	// 3. Call LLM with group keys (batched to stay within token limits)
	mapping := make(map[string]normMapping)
	batchSize := 50
	for i := 0; i < len(keys); i += batchSize {
		end := i + batchSize
		if end > len(keys) {
			end = len(keys)
		}
		batch := keys[i:end]
		reportProgress(end, len(keys))
		log.Printf("normalize: LLM batch %d-%d of %d", i+1, end, len(keys))

		batchMap, err := callNormalizeLLM(batch)
		if err != nil {
			return fmt.Errorf("LLM batch %d-%d: %w", i+1, end, err)
		}
		for k, v := range batchMap {
			mapping[k] = v
		}
	}

	// 4. Expand: each original name in a group gets the group's canonical mapping
	fullMapping := make(map[string]normMapping)
	for key, origNames := range groups {
		if m, ok := mapping[key]; ok {
			for _, orig := range origNames {
				fullMapping[orig] = m
			}
		}
	}

	log.Printf("normalize: LLM mapped %d groups → %d original names covered", len(mapping), len(fullMapping))

	// 5. Save LabTest entries for any new LOINC codes
	seenLoinc := make(map[string]bool)
	var labTests []LabTest
	for _, m := range fullMapping {
		if m.Loinc == "" || seenLoinc[m.Loinc] {
			continue
		}
		seenLoinc[m.Loinc] = true
		dir := m.Direction
		if dir == "" {
			dir = DirRange
		}
		factor := m.SIFactor
		if factor == 0 {
			factor = 1.0
		}
		labTests = append(labTests, LabTest{
			LoincID:   m.Loinc,
			Name:      m.Name,
			SIUnit:    m.SIUnit,
			Direction: dir,
			SIFactor:  ToLabScale(factor),
		})
	}
	for _, t := range labTests {
		RefExec(`INSERT OR IGNORE INTO lab_test (loinc_id, name, si_unit, direction, si_factor) VALUES (?, ?, ?, ?, ?)`,
			t.LoincID, t.Name, t.SIUnit, t.Direction, t.SIFactor)
	}
	if len(labTests) > 0 {
		log.Printf("normalize: saved %d lab tests", len(labTests))
	}

	// 5. Apply mapping to loaded entries, save only changed ones
	var toSave []Entry
	for _, e := range entries {
		if e.ParentID == "" {
			continue
		}
		norm, ok := fullMapping[e.Type]
		if !ok {
			continue
		}

		var data map[string]interface{}
		if json.Unmarshal([]byte(e.Data), &data) != nil {
			data = make(map[string]interface{})
		}

		// Skip if already fully normalized
		existingName, _ := data["normalized_name"].(string)
		existingLoinc, _ := data["loinc"].(string)
		needsSearchKey := (norm.Loinc != "" && e.SearchKey == "")
		needsSearchKey2 := e.SearchKey2 == ""
		if existingName == norm.Name && (norm.Loinc == "" || existingLoinc == norm.Loinc) && !needsSearchKey && !needsSearchKey2 {
			continue
		}

		data["normalized_name"] = norm.Name
		data["abbreviation"] = norm.Abbr
		if norm.Loinc != "" {
			data["loinc"] = norm.Loinc
		}
		if norm.SIUnit != "" {
			data["si_unit"] = norm.SIUnit
		}
		if norm.SIFactor != 0 && norm.SIFactor != 1.0 {
			data["si_factor"] = norm.SIFactor
		}
		b, _ := json.Marshal(data)
		e.Data = string(b)

		// Update SearchKey with LOINC code, SearchKey2 with canonical test name
		if norm.Loinc != "" {
			e.SearchKey = norm.Loinc
		}
		e.SearchKey2 = strings.ToLower(norm.Name)

		// Rebuild Summary: "Abbr: value unit"
		unit, _ := data["unit"].(string)
		summary := norm.Abbr + ": " + e.Value
		if unit != "" {
			summary += " " + unit
		}
		e.Summary = summary

		toSave = append(toSave, *e)
	}

	if len(toSave) == 0 {
		log.Printf("normalize: no changes needed")
		return nil
	}

	log.Printf("normalize: updating %d entries", len(toSave))
	ptrs := make([]*Entry, len(toSave))
	for i := range toSave {
		ptrs[i] = &toSave[i]
	}
	return EntryWrite("", ptrs...)
}

// normalizeKey reduces a test name to a heuristic grouping key.
// Groups obvious duplicates: POCT variants, specimen suffixes, case.
func normalizeKey(name string) string {
	s := strings.ToLower(strings.TrimSpace(name))
	s = strings.TrimPrefix(s, "poct ")
	// Strip specimen-type suffixes only (not qualifiers like ", total", ", direct")
	for _, suf := range []string{", whole blood", ", wblood", ", wb", ", wbl", ", blood", ", s/p", " ach"} {
		s = strings.TrimSuffix(s, suf)
	}
	return strings.TrimSpace(s)
}

type normMapping struct {
	Name      string  `json:"name"`
	Abbr      string  `json:"abbr"`
	Loinc     string  `json:"loinc"`
	SIUnit    string  `json:"si_unit"`
	SIFactor  float64 `json:"si_factor"`
	Direction string  `json:"direction"`
}

func callNormalizeLLM(names []string) (map[string]normMapping, error) {
	nameList := strings.Join(names, "\n")

	prompt := fmt.Sprintf(`Given these medical test names from a single patient's records, normalize each to a canonical name, abbreviation, LOINC code, SI unit, conversion factor, and direction.

Rules:
- Use standard medical abbreviations: WBC, RBC, Hgb, Hct, PLT, Na, K, Cl, CO2, BUN, Cr, Ca, Glu, ALT, AST, ALP, Bili, Alb, TP, Mg, Phos, Fe, etc.
- For tests without standard abbreviations, use a short canonical name as abbreviation
- Keep abbreviations concise (1-8 chars)
- If two names are the same test, give them the same canonical name and abbreviation
- loinc: the most common LOINC code for this test (e.g. "718-7" for Hemoglobin). Use "" if unknown.
- si_unit: the standard SI unit (e.g. "g/L", "mmol/L", "10^9/L"). Use "" if not numeric.
- si_factor: multiplier to convert from the most common conventional unit to SI. E.g. Hemoglobin g/dL→g/L = 10.0. Use 1.0 if already SI or unknown.
- direction: "range" if both high and low are bad (most tests), "lower_better" if low values are healthy (CRP, LDL, triglycerides, glucose), "higher_better" if high values are healthy (HDL). Default to "range".

Return a JSON object where each key is the EXACT input name, value is {"name":"Canonical Name","abbr":"Abbreviation","loinc":"CODE","si_unit":"unit","si_factor":1.0,"direction":"range"}.

Test names:
%s`, nameList)

	maxTokens := 32768
	temp := 0.0
	model := "gemini-3.1-pro-preview"
	config := &GeminiConfig{
		Temperature:     &temp,
		MaxOutputTokens: &maxTokens,
		Model:           &model,
	}

	resp, err := CallGeminiMultimodal([]GeminiPart{{Text: prompt}}, config)
	if err != nil {
		return nil, err
	}

	// Gemini sometimes returns object, sometimes array of objects
	var mapping map[string]normMapping
	if err := json.Unmarshal([]byte(resp), &mapping); err != nil {
		var arr []map[string]normMapping
		if err2 := json.Unmarshal([]byte(resp), &arr); err2 != nil {
			return nil, fmt.Errorf("parse response: %w (first 300 chars: %.300s)", err, resp)
		}
		mapping = make(map[string]normMapping)
		for _, item := range arr {
			for k, v := range item {
				mapping[k] = v
			}
		}
	}

	return mapping, nil
}