package main

import (
	"fmt"
	"os"
	"sort"
	"strings"

	sitter "github.com/tree-sitter/go-tree-sitter"
	ts_go "github.com/tree-sitter/tree-sitter-go/bindings/go"
)

// Sample Go source to highlight
const source = `
	package main 

	func main () {
		println("Hello" + 5)
	}
`

type Highlight struct {
	StartRow uint // 0-indexed line number
	StartCol uint // 0-indexed column (bytes)
	EndRow   uint
	EndCol   uint
	Capture  string
}

// Theme maps capture names to ANSI escape codes.
// In your editor you'd use lipgloss styles instead.
var theme = map[string]string{
	"keyword":               "\033[1;35m", // bold magenta
	"keyword.type":          "\033[1;35m",
	"keyword.function":      "\033[1;35m",
	"keyword.return":        "\033[1;35m",
	"keyword.coroutine":     "\033[1;35m",
	"keyword.repeat":        "\033[1;35m",
	"keyword.import":        "\033[1;35m",
	"keyword.conditional":   "\033[1;35m",
	"type":                  "\033[33m", // yellow
	"type.builtin":          "\033[33m",
	"type.definition":       "\033[1;33m", // bold yellow
	"function":              "\033[1;34m", // bold blue
	"function.call":         "\033[34m",   // blue
	"function.method":       "\033[34m",
	"function.method.call":  "\033[34m",
	"function.builtin":      "\033[1;31m",
	"variable":              "\033[37m",   // white
	"variable.parameter":    "\033[3;37m", // italic white
	"variable.member":       "\033[37m",
	"constant":              "\033[1;36m", // bold cyan
	"constant.builtin":      "\033[1;36m",
	"string":                "\033[32m", // green
	"string.escape":         "\033[1;32m",
	"number":                "\033[36m", // cyan
	"number.float":          "\033[36m",
	"boolean":               "\033[36m",
	"operator":              "\033[93m",   // bright yellow
	"comment":               "\033[2;37m", // dim
	"comment.documentation": "\033[2;37m",
	"module":                "\033[35m", // magenta
	"label":                 "\033[33m",
	"property":              "\033[37m",
	"constructor":           "\033[1;33m",
	"punctuation.delimiter": "\033[37m",
	"punctuation.bracket":   "\033[37m",
}

const reset = "\033[0m"

func main() {
	code := []byte(source)
	lines := strings.Split(source, "\n")

	// --- Step 1: Parse ---
	lang := sitter.NewLanguage(ts_go.Language())
	parser := sitter.NewParser()
	defer parser.Close()
	parser.SetLanguage(lang)

	tree := parser.Parse(code, nil)
	defer tree.Close()
	root := tree.RootNode()

	// --- Step 2: Load query from highlights.scm ---
	queryBytes, err := os.ReadFile("queries/go/highlights.scm")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to read highlights.scm: %v\n", err)
		return
	}

	query, queryErr := sitter.NewQuery(lang, string(queryBytes))
	if queryErr != nil {
		fmt.Fprintf(os.Stderr, "Query error: %v\n", queryErr)
		return
	}
	defer query.Close()

	// --- Step 3: Run query ---
	cursor := sitter.NewQueryCursor()
	defer cursor.Close()
	captures := cursor.Captures(query, root, code)

	var highlights []Highlight
	for match, captureIdx := captures.Next(); match != nil; match, captureIdx = captures.Next() {
		capture := match.Captures[captureIdx]
		captureName := query.CaptureNames()[capture.Index]
		// Skip @spell — it's a nvim spellcheck hint, not a highlight
		if captureName == "spell" {
			continue
		}
		node := capture.Node
		start := node.StartPosition()
		end := node.EndPosition()
		highlights = append(highlights, Highlight{
			StartRow: start.Row,
			StartCol: start.Column,
			EndRow:   end.Row,
			EndCol:   end.Column,
			Capture:  captureName,
		})
	}

	// --- Step 4: Show captures with positions ---
	fmt.Println("=== Captures (row:col → row:col) ===")
	for _, h := range highlights {
		// Extract text for display using the source lines
		text := extractText(lines, h)
		fmt.Printf("  %d:%-2d → %d:%-2d  @%-22s %q\n",
			h.StartRow, h.StartCol, h.EndRow, h.EndCol, h.Capture, text)
	}
	fmt.Println()

	// --- Step 5: Render with colors using row:col positions ---
	// Build a per-line map of column ranges to capture names.
	// Sort so wider (less specific) ranges come first — last writer wins.
	sort.Slice(highlights, func(i, j int) bool {
		if highlights[i].StartRow == highlights[j].StartRow {
			if highlights[i].StartCol == highlights[j].StartCol {
				// Wider range first so more specific overwrites it
				if highlights[i].EndRow == highlights[j].EndRow {
					return highlights[i].EndCol > highlights[j].EndCol
				}
				return highlights[i].EndRow > highlights[j].EndRow
			}
			return highlights[i].StartCol < highlights[j].StartCol
		}
		return highlights[i].StartRow < highlights[j].StartRow
	})

	// captureAt[row][col] = capture name (last writer wins)
	captureAt := make(map[uint]map[uint]string)
	for _, h := range highlights {
		for row := h.StartRow; row <= h.EndRow; row++ {
			if captureAt[row] == nil {
				captureAt[row] = make(map[uint]string)
			}
			startCol := uint(0)
			if row == h.StartRow {
				startCol = h.StartCol
			}
			endCol := uint(len(lines[row]))
			if row == h.EndRow {
				endCol = h.EndCol
			}
			for col := startCol; col < endCol; col++ {
				captureAt[row][col] = h.Capture
			}
		}
	}

	fmt.Println("=== Colored output ===")
	printColored(lines, captureAt)

	// =====================================================================
	// INCREMENTAL PARSING DEMO
	// =====================================================================
	// When a user types in your editor, you don't re-parse the whole file.
	// Instead you:
	//   1. Tell the OLD tree what changed (tree.Edit)
	//   2. Parse the new source, passing the old tree
	//   3. Tree-sitter reuses unchanged nodes and only re-parses the edit
	//   4. Use ChangedRanges to know which lines need re-highlighting
	//
	// This is O(edit size + log(file size)) instead of O(file size).

	fmt.Println("\n========================================")
	fmt.Println("=== INCREMENTAL PARSE DEMO ===")
	fmt.Println("========================================")
	fmt.Println()

	// Simulate: user changes "Hello" → "Goodbye" on row 4
	// Before: println("Hello" + 5)
	// After:  println("Goodbye" + 5)
	oldSource := source
	newSource := strings.Replace(oldSource, `"Hello"`, `"Goodbye"`, 1)

	// Find where the edit happened (in a real editor you already know this
	// from the keystroke — you don't need to search for it)
	editStart := strings.Index(oldSource, `"Hello"`)
	oldEnd := editStart + len(`"Hello"`)
	newEnd := editStart + len(`"Goodbye"`)

	// Convert byte offset to row:col for the InputEdit
	editStartPoint := byteToPoint(oldSource, uint(editStart))
	oldEndPoint := byteToPoint(oldSource, uint(oldEnd))
	newEndPoint := byteToPoint(newSource, uint(newEnd))

	fmt.Printf("Edit: replaced %q → %q\n", "Hello", "Goodbye")
	fmt.Printf("  at byte %d, row %d col %d\n", editStart, editStartPoint.Row, editStartPoint.Column)
	fmt.Println()

	// Step 1: Tell the old tree what changed
	tree.Edit(&sitter.InputEdit{
		StartByte:      uint(editStart),
		OldEndByte:     uint(oldEnd),
		NewEndByte:     uint(newEnd),
		StartPosition:  editStartPoint,
		OldEndPosition: oldEndPoint,
		NewEndPosition: newEndPoint,
	})

	// Step 2: Parse the new source, passing the old (edited) tree.
	// Tree-sitter will REUSE all nodes that weren't affected by the edit
	// and only re-parse the region around the change.
	newCode := []byte(newSource)
	newTree := parser.Parse(newCode, tree)
	defer newTree.Close()

	// Step 3: See exactly which ranges changed
	changedRanges := newTree.ChangedRanges(tree)
	fmt.Printf("Changed ranges: %d\n", len(changedRanges))
	for i, r := range changedRanges {
		fmt.Printf("  range %d: row %d:%d → row %d:%d\n",
			i, r.StartPoint.Row, r.StartPoint.Column, r.EndPoint.Row, r.EndPoint.Column)
	}
	fmt.Println()

	// In your editor, you'd ONLY re-run the highlight query on the changed
	// ranges (using cursor.SetByteRange or cursor.SetPointRange), then
	// update just those lines in your display. Everything else stays cached.

	// For this demo, let's re-highlight the full new tree to show the result
	newRoot := newTree.RootNode()
	newLines := strings.Split(newSource, "\n")

	cursor2 := sitter.NewQueryCursor()
	defer cursor2.Close()
	newCaptures := cursor2.Captures(query, newRoot, newCode)

	var newHighlights []Highlight
	for match, captureIdx := newCaptures.Next(); match != nil; match, captureIdx = newCaptures.Next() {
		capture := match.Captures[captureIdx]
		captureName := query.CaptureNames()[capture.Index]
		if captureName == "spell" {
			continue
		}
		node := capture.Node
		start := node.StartPosition()
		end := node.EndPosition()
		newHighlights = append(newHighlights, Highlight{
			StartRow: start.Row, StartCol: start.Column,
			EndRow: end.Row, EndCol: end.Column,
			Capture: captureName,
		})
	}

	sort.Slice(newHighlights, func(i, j int) bool {
		if newHighlights[i].StartRow == newHighlights[j].StartRow {
			if newHighlights[i].StartCol == newHighlights[j].StartCol {
				if newHighlights[i].EndRow == newHighlights[j].EndRow {
					return newHighlights[i].EndCol > newHighlights[j].EndCol
				}
				return newHighlights[i].EndRow > newHighlights[j].EndRow
			}
			return newHighlights[i].StartCol < newHighlights[j].StartCol
		}
		return newHighlights[i].StartRow < newHighlights[j].StartRow
	})

	newCaptureAt := make(map[uint]map[uint]string)
	for _, h := range newHighlights {
		for row := h.StartRow; row <= h.EndRow; row++ {
			if newCaptureAt[row] == nil {
				newCaptureAt[row] = make(map[uint]string)
			}
			startCol := uint(0)
			if row == h.StartRow {
				startCol = h.StartCol
			}
			endCol := uint(len(newLines[row]))
			if row == h.EndRow {
				endCol = h.EndCol
			}
			for col := startCol; col < endCol; col++ {
				newCaptureAt[row][col] = h.Capture
			}
		}
	}

	fmt.Println("=== After edit (colored output) ===")
	printColored(newLines, newCaptureAt)
}

// printColored renders source lines with ANSI colors based on the capture map.
func printColored(lines []string, captureAt map[uint]map[uint]string) {
	for row, line := range lines {
		currentCapture := ""
		for col := uint(0); col < uint(len(line)); col++ {
			cap := ""
			if rowMap, ok := captureAt[uint(row)]; ok {
				cap = rowMap[col]
			}
			if cap != currentCapture {
				if currentCapture != "" {
					fmt.Print(reset)
				}
				if color, ok := theme[cap]; ok {
					fmt.Print(color)
				}
				currentCapture = cap
			}
			fmt.Print(string(line[col]))
		}
		if currentCapture != "" {
			fmt.Print(reset)
		}
		fmt.Println()
	}
}

// byteToPoint converts a byte offset into a row:col Point.
func byteToPoint(src string, offset uint) sitter.Point {
	row := uint(0)
	col := uint(0)
	for i := range offset {
		if src[i] == '\n' {
			row++
			col = 0
		} else {
			col++
		}
	}
	return sitter.NewPoint(row, col)
}

// extractText pulls the highlighted text from source lines using row:col positions.
func extractText(lines []string, h Highlight) string {
	if h.StartRow == h.EndRow {
		line := lines[h.StartRow]
		end := min(h.EndCol, uint(len(line)))
		return line[h.StartCol:end]
	}
	// Multi-line highlight (rare, but possible for block comments etc.)
	var result string
	for row := h.StartRow; row <= h.EndRow; row++ {
		line := lines[row]
		start := uint(0)
		if row == h.StartRow {
			start = h.StartCol
		}
		end := uint(len(line))
		if row == h.EndRow {
			end = h.EndCol
		}
		if result != "" {
			result += "\n"
		}
		result += line[start:end]
	}
	return result
}