package cartography

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

// Atlas is the bidirectional English↔Go map of the entire codebase.
type Atlas struct {
	// All entries, keyed by ID ("package.Name" or "package.Type.Method").
	Entries map[string]*Entry `json:"entries"`

	// Metadata.
	Generation     int     `json:"generation"`
	TotalEntries   int     `json:"total_entries"`
	MeanConfidence float64 `json:"mean_confidence"`
	OracleCalls    int     `json:"oracle_calls"` // lifetime

	// In-memory indexes (rebuilt on Load, not serialized).
	byPackage map[string][]*Entry
	byConcept map[string][]*Entry
	byFile    map[string][]*Entry
}

// NewAtlas creates an empty atlas.
func NewAtlas() *Atlas {
	return &Atlas{
		Entries:   make(map[string]*Entry),
		byPackage: make(map[string][]*Entry),
		byConcept: make(map[string][]*Entry),
		byFile:    make(map[string][]*Entry),
	}
}

// Add inserts or replaces an entry and updates metadata.
func (a *Atlas) Add(e *Entry) {
	a.Entries[e.ID] = e
	a.updateStats()
}

// BuildIndexes constructs the in-memory bidirectional indexes from entries.
func (a *Atlas) BuildIndexes() {
	a.byPackage = make(map[string][]*Entry)
	a.byConcept = make(map[string][]*Entry)
	a.byFile = make(map[string][]*Entry)

	for _, e := range a.Entries {
		a.byPackage[e.Package] = append(a.byPackage[e.Package], e)
		a.byFile[e.FilePath] = append(a.byFile[e.FilePath], e)
		for _, c := range e.Concepts {
			a.byConcept[c] = append(a.byConcept[c], e)
		}
	}
}

// LookupGo returns the entry for a Go declaration.
// Accepts "package.Name", "package.Type.Method", or just "Name" (searches all).
func (a *Atlas) LookupGo(id string) *Entry {
	// Direct match.
	if e, ok := a.Entries[id]; ok {
		return e
	}
	// Search by bare name.
	for _, e := range a.Entries {
		if e.Name == id {
			return e
		}
	}
	return nil
}

// LookupEnglish returns Go declarations matching an English query.
// Searches concept tags, descriptions, and doc comments.
// Returns entries sorted by relevance.
func (a *Atlas) LookupEnglish(query string) []*Entry {
	words := splitQuery(query)
	if len(words) == 0 {
		return nil
	}

	type scored struct {
		entry *Entry
		score float64
	}
	var results []scored

	for _, e := range a.Entries {
		score := 0.0
		for _, w := range words {
			wl := strings.ToLower(w)
			// Concept match (strongest signal).
			for _, c := range e.Concepts {
				if c == wl {
					score += 3.0
				}
			}
			// Name match.
			if strings.EqualFold(e.Name, w) {
				score += 5.0
			}
			// Package match.
			if strings.EqualFold(e.Package, w) {
				score += 2.0
			}
			// Description match.
			if strings.Contains(strings.ToLower(e.Description), wl) {
				score += 1.0
			}
			// Doc comment match.
			if strings.Contains(strings.ToLower(e.DocComment), wl) {
				score += 0.5
			}
		}
		if score > 0 {
			results = append(results, scored{e, score})
		}
	}

	sort.Slice(results, func(i, j int) bool {
		return results[i].score > results[j].score
	})

	entries := make([]*Entry, len(results))
	for i, r := range results {
		entries[i] = r.entry
	}
	return entries
}

// LookupPackage returns all entries in a package.
func (a *Atlas) LookupPackage(pkg string) []*Entry {
	return a.byPackage[pkg]
}

// Packages returns all unique package names in the atlas.
func (a *Atlas) Packages() []string {
	pkgs := make([]string, 0, len(a.byPackage))
	for p := range a.byPackage {
		pkgs = append(pkgs, p)
	}
	sort.Strings(pkgs)
	return pkgs
}

// Summary returns a compact text summary of the atlas suitable for oracle
// context. maxChars limits the output size (0 = unlimited).
// This produces richer context than bare type/function name lists.
func (a *Atlas) Summary(maxChars int) string {
	var b strings.Builder

	pkgs := a.Packages()
	for _, pkg := range pkgs {
		entries := a.byPackage[pkg]
		if len(entries) == 0 {
			continue
		}

		// Package header.
		fmt.Fprintf(&b, "Package %s", pkg)
		// Find a package-level description from the highest-confidence entry.
		bestDesc := ""
		for _, e := range entries {
			if e.DocComment != "" && e.Confidence >= 0.3 {
				bestDesc = firstSentence(e.DocComment)
				break
			}
		}
		if bestDesc != "" {
			fmt.Fprintf(&b, " — %s", bestDesc)
		}
		b.WriteString("\n")

		// Sort entries: types first, then funcs, then methods.
		sorted := make([]*Entry, len(entries))
		copy(sorted, entries)
		sort.Slice(sorted, func(i, j int) bool {
			return kindOrder(sorted[i].Kind) < kindOrder(sorted[j].Kind)
		})

		for _, e := range sorted {
			if !e.Exported {
				continue // only include exported in summary
			}
			sig := e.Signature
			if sig == "" {
				sig = e.Name
			}
			desc := firstSentence(e.Description)
			if desc != "" {
				fmt.Fprintf(&b, "  %s — %s\n", sig, desc)
			} else {
				fmt.Fprintf(&b, "  %s\n", sig)
			}

			if maxChars > 0 && b.Len() > maxChars {
				fmt.Fprintf(&b, "  ... (%d more entries)\n", a.TotalEntries-len(a.Entries))
				return b.String()
			}
		}
	}

	return b.String()
}

// Merge combines a freshly-extracted atlas with this one.
// New entries are added. Existing entries keep oracle-enriched descriptions
// but update mechanical fields (signature, params, returns) from fresh.
// Entries that no longer exist in fresh are removed.
func (a *Atlas) Merge(fresh *Atlas) {
	// Update existing + add new.
	for id, freshEntry := range fresh.Entries {
		if existing, ok := a.Entries[id]; ok {
			// Keep oracle-enriched fields.
			if existing.Source == "oracle" || existing.Source == "self-revised" {
				freshEntry.Description = existing.Description
				freshEntry.Contract = existing.Contract
				freshEntry.EdgeCases = existing.EdgeCases
				freshEntry.Confidence = existing.Confidence
				freshEntry.Source = existing.Source
				freshEntry.Revision = existing.Revision
				// Keep oracle-enriched param semantics.
				for i := range freshEntry.Params {
					if i < len(existing.Params) && existing.Params[i].Semantic != "" {
						freshEntry.Params[i].Semantic = existing.Params[i].Semantic
					}
				}
				for i := range freshEntry.Returns {
					if i < len(existing.Returns) && existing.Returns[i].Semantic != "" {
						freshEntry.Returns[i].Semantic = existing.Returns[i].Semantic
					}
				}
			}
			// Keep test data if fresh doesn't have it.
			if !freshEntry.HasTest && existing.HasTest {
				freshEntry.HasTest = existing.HasTest
				freshEntry.TestFuncs = existing.TestFuncs
				freshEntry.TestFile = existing.TestFile
				freshEntry.TestPatterns = existing.TestPatterns
				freshEntry.Validation = existing.Validation
			}
		}
		a.Entries[id] = freshEntry
	}

	// Remove entries not in fresh (code was deleted).
	for id := range a.Entries {
		if _, ok := fresh.Entries[id]; !ok {
			delete(a.Entries, id)
		}
	}

	a.updateStats()
	a.BuildIndexes()
}

// Save writes the atlas to a JSON file.
func (a *Atlas) Save(path string) error {
	a.updateStats()
	data, err := json.MarshalIndent(a, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, data, 0o644)
}

// Load reads an atlas from a JSON file and rebuilds indexes.
func Load(path string) (*Atlas, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	a := &Atlas{}
	if err := json.Unmarshal(data, a); err != nil {
		return nil, err
	}
	if a.Entries == nil {
		a.Entries = make(map[string]*Entry)
	}
	a.BuildIndexes()
	return a, nil
}

// LowestConfidence returns the n entries with the lowest confidence,
// sorted ascending. Ties are broken by ID for deterministic ordering.
func (a *Atlas) LowestConfidence(n int) []*Entry {
	all := make([]*Entry, 0, len(a.Entries))
	for _, e := range a.Entries {
		all = append(all, e)
	}
	sort.Slice(all, func(i, j int) bool {
		if all[i].Confidence != all[j].Confidence {
			return all[i].Confidence < all[j].Confidence
		}
		return all[i].ID < all[j].ID
	})
	if n > len(all) {
		n = len(all)
	}
	return all[:n]
}

func (a *Atlas) updateStats() {
	a.TotalEntries = len(a.Entries)
	if a.TotalEntries == 0 {
		a.MeanConfidence = 0
		return
	}
	total := 0.0
	for _, e := range a.Entries {
		total += e.Confidence
	}
	a.MeanConfidence = total / float64(a.TotalEntries)
}

// splitQuery breaks an English query into searchable words.
func splitQuery(query string) []string {
	words := strings.Fields(strings.ToLower(query))
	// Filter stop words.
	var result []string
	for _, w := range words {
		if len(w) >= 3 && !stopWords[w] {
			result = append(result, w)
		}
	}
	return result
}

var stopWords = map[string]bool{
	"the": true, "and": true, "for": true, "are": true,
	"but": true, "not": true, "you": true, "all": true,
	"can": true, "had": true, "her": true, "was": true,
	"one": true, "our": true, "out": true, "has": true,
	"its": true, "that": true, "with": true, "this": true,
	"from": true, "they": true, "been": true, "have": true,
	"which": true, "when": true, "what": true, "does": true,
	"how": true, "where": true,
}

// firstSentence returns the first sentence of a string (up to first period
// or newline), trimmed.
func firstSentence(s string) string {
	s = strings.TrimSpace(s)
	if i := strings.IndexByte(s, '.'); i >= 0 && i < 120 {
		return s[:i+1]
	}
	if i := strings.IndexByte(s, '\n'); i >= 0 {
		return strings.TrimSpace(s[:i])
	}
	if len(s) > 120 {
		return s[:120] + "..."
	}
	return s
}

func kindOrder(kind string) int {
	switch kind {
	case "type", "interface":
		return 0
	case "func":
		return 1
	case "method":
		return 2
	default:
		return 3
	}
}