package main

import (
	"fmt"
	"go/ast"
	"go/token"
	"strconv"
	"strings"
)

// KindInfo holds information about a Kind constant
type KindInfo struct {
	Name    string
	Value   int64
	Comment string
}

// KindProcessor processes Kind enum definitions
type KindProcessor struct {
	kinds []KindInfo
}

// NewKindProcessor creates a new kind processor
func NewKindProcessor() *KindProcessor {
	return &KindProcessor{
		kinds: make([]KindInfo, 0),
	}
}

// Process extracts Kind information from Go AST files
func (p *KindProcessor) Process(files map[string]*ast.File) error {
	// Look for the kind.go file
	for filename, file := range files {
		if strings.Contains(filename, "kind.go") {
			p.processFile(file)
		}
	}
	return nil
}

// processFile processes the kind.go file
func (p *KindProcessor) processFile(file *ast.File) {
	// Find the Kind type
	var kindType bool
	ast.Inspect(file, func(n ast.Node) bool {
		if typeSpec, ok := n.(*ast.TypeSpec); ok {
			if typeSpec.Name.Name == "Kind" {
				kindType = true
			}
		}
		return true
	})

	if !kindType {
		return
	}

	// Find Kind constants
	ast.Inspect(file, func(n ast.Node) bool {
		if genDecl, ok := n.(*ast.GenDecl); ok && genDecl.Tok == token.CONST {
			p.processConstDecl(genDecl)
		}
		return true
	})
}

// processConstDecl processes a const declaration for Kind values
func (p *KindProcessor) processConstDecl(decl *ast.GenDecl) {
	var isKindConst bool
	var lastValue int64 = -1

	for _, spec := range decl.Specs {
		valueSpec, ok := spec.(*ast.ValueSpec)
		if !ok {
			continue
		}

		// Check if this is a Kind constant
		if valueSpec.Type != nil {
			if ident, ok := valueSpec.Type.(*ast.Ident); ok && ident.Name == "Kind" {
				isKindConst = true
				lastValue = -1 // Reset for new const block
			} else {
				isKindConst = false
			}
		}

		if !isKindConst {
			continue
		}

		// Process each constant
		for i, name := range valueSpec.Names {
			// Skip if name starts with "Kind" prefix (it's the enum name)
			nameStr := name.Name
			if !strings.HasPrefix(nameStr, "Kind") {
				continue
			}

			// Skip if this is an alias (assigned to another Kind identifier)
			if i < len(valueSpec.Values) {
				if p.isAlias(valueSpec.Values[i]) {
					continue
				}
			}

			// Regular Kind constant
			kind := KindInfo{
				Name: strings.TrimPrefix(nameStr, "Kind"),
			}

			// Extract value
			if i < len(valueSpec.Values) {
				kind.Value = p.extractValue(valueSpec.Values[i], lastValue)
			} else {
				kind.Value = lastValue + 1
			}
			lastValue = kind.Value

			// Extract comment
			if valueSpec.Comment != nil {
				kind.Comment = strings.TrimSpace(valueSpec.Comment.Text())
			}

			p.kinds = append(p.kinds, kind)
		}
	}
}

// isAlias checks if an expression is an alias to another Kind constant
func (p *KindProcessor) isAlias(expr ast.Expr) bool {
	switch e := expr.(type) {
	case *ast.Ident:
		if strings.HasPrefix(e.Name, "Kind") && e.Name != "iota" {
			return true
		}
	case *ast.BinaryExpr:
		return true
	}
	return false
}

// extractValue extracts the numeric value from an expression
func (p *KindProcessor) extractValue(expr ast.Expr, lastValue int64) int64 {
	switch e := expr.(type) {
	case *ast.BasicLit:
		if e.Kind == token.INT {
			val, _ := strconv.ParseInt(e.Value, 0, 64)
			return val
		}
	case *ast.Ident:
		if e.Name == "iota" {
			return lastValue + 1
		}
	}
	return lastValue + 1
}

// GenerateRust generates Rust code for updating the SyntaxKind enum
func (p *KindProcessor) GenerateRust() string {
	var output strings.Builder

	output.WriteString("// Generated by ast-gen - DO NOT EDIT\n")
	output.WriteString("use strum::{Display, EnumCount, FromRepr};\n\n")
	output.WriteString("#[repr(u16)]\n")
	output.WriteString("#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Display, FromRepr, EnumCount)]\n")
	output.WriteString("pub enum SyntaxKind {\n")

	// Generate enum variants
	for _, kind := range p.kinds {
		variantName := kind.Name

		// Special handling for certain names
		switch variantName {
		case "EndOfFileToken":
			variantName = "EndOfFile"
		case "NotEmittedTypeElement":
			continue
		case "Count":
			continue
		}

		if kind.Comment != "" {
			output.WriteString(fmt.Sprintf("    /// %s\n", kind.Comment))
		}

		if kind.Value == 0 {
			output.WriteString(fmt.Sprintf("    %s = 0x%02x,\n", variantName, kind.Value))
		} else {
			output.WriteString(fmt.Sprintf("    %s,\n", variantName))
		}
	}

	output.WriteString("}\n")

	return output.String()
}

// OutputFile returns the output file name
func (p *KindProcessor) OutputFile() string {
	return "syntax_kind.rs"
}
