// Build from project root:
// go build -o tools/parse_go_ast tools/parse_go_ast.go

package main

import (
	"encoding/json"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"strings"
)

type GoDecl struct {
	Name       string   `json:"name"`
	Type       string   `json:"type"`               // "func", "method", "struct", "interface", "type", "var", "const"
	Receiver   string   `json:"receiver,omitempty"` // e.g., "*User"
	FullName   string   `json:"full_name,omitempty"`
	Fields     []string `json:"fields,omitempty"`     // for structs
	Methods    []string `json:"methods,omitempty"`    // for interfaces
	Parameters []string `json:"parameters,omitempty"` // for functions/methods
	ReturnType string   `json:"return_type,omitempty"`
	DocComment string   `json:"doc_comment,omitempty"`
	IsExported bool     `json:"is_exported,omitempty"` // starts with capital letter
	StartLine  int      `json:"start_line"`
	EndLine    int      `json:"end_line"`
}

func astTypeToString(expr ast.Expr) string {
	if expr == nil {
		return ""
	}
	switch t := expr.(type) {
	case *ast.Ident:
		return t.Name
	case *ast.StarExpr:
		return "*" + astTypeToString(t.X)
	case *ast.SelectorExpr:
		return astTypeToString(t.X) + "." + t.Sel.Name
	case *ast.ArrayType:
		return "[]" + astTypeToString(t.Elt)
	case *ast.MapType:
		return "map[" + astTypeToString(t.Key) + "]" + astTypeToString(t.Value)
	case *ast.InterfaceType:
		return "interface{}"
	case *ast.StructType:
		return "struct{}"
	case *ast.FuncType:
		return "func"
	case *ast.ChanType:
		return "chan " + astTypeToString(t.Value)
	case *ast.Ellipsis:
		return "..." + astTypeToString(t.Elt)
	default:
		return "unknown"
	}
}

func extractDocComment(comments []*ast.CommentGroup, pos token.Pos, fset *token.FileSet) string {
	if len(comments) == 0 {
		return ""
	}
	line := fset.Position(pos).Line

	// Look for comments immediately before the declaration
	for i := len(comments) - 1; i >= 0; i-- {
		cg := comments[i]
		cgEndLine := fset.Position(cg.End()).Line
		cgStartLine := fset.Position(cg.Pos()).Line

		// Comment should be within 1 line of the declaration
		if cgEndLine < line && line-cgEndLine <= 1 {
			text := strings.TrimSpace(cg.Text())
			// Clean up comment markers
			text = strings.ReplaceAll(text, "/*", "")
			text = strings.ReplaceAll(text, "*/", "")
			return strings.TrimSpace(text)
		}

		// If we've gone too far back, stop looking
		if cgStartLine < line-10 {
			break
		}
	}
	return ""
}

func extractParameters(fields *ast.FieldList) []string {
	if fields == nil {
		return nil
	}

	var params []string
	for _, field := range fields.List {
		typeStr := astTypeToString(field.Type)
		if len(field.Names) == 0 {
			// Unnamed parameter
			params = append(params, typeStr)
		} else {
			for _, name := range field.Names {
				params = append(params, name.Name+": "+typeStr)
			}
		}
	}
	return params
}

func extractReturnType(results *ast.FieldList) string {
	if results == nil || len(results.List) == 0 {
		return ""
	}

	if len(results.List) == 1 {
		return astTypeToString(results.List[0].Type)
	}

	// Multiple return values
	var types []string
	for _, field := range results.List {
		types = append(types, astTypeToString(field.Type))
	}
	return "(" + strings.Join(types, ", ") + ")"
}

func isExported(name string) bool {
	if len(name) == 0 {
		return false
	}
	firstRune := []rune(name)[0]
	return firstRune >= 'A' && firstRune <= 'Z'
}

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintln(os.Stderr, "Usage: parse_go_ast <file.go>")
		os.Exit(1)
	}

	filename := os.Args[1]
	src, err := os.ReadFile(filename)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERR: Failed to read file: %v\n", err)
		os.Exit(2)
	}

	fset := token.NewFileSet()
	fileNode, err := parser.ParseFile(fset, filename, src, parser.ParseComments)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERR: Failed to parse Go file: %v\n", err)
		os.Exit(3)
	}

	var decls []GoDecl

	// Process declarations
	for _, d := range fileNode.Decls {
		switch d := d.(type) {
		case *ast.FuncDecl:
			name := d.Name.Name
			recvType := ""

			if d.Recv != nil && len(d.Recv.List) > 0 {
				recvType = astTypeToString(d.Recv.List[0].Type)
			}

			typ := "function"
			fullName := name
			if recvType != "" {
				typ = "method"
				fullName = fmt.Sprintf("(%s).%s", recvType, name)
			}

			doc := extractDocComment(fileNode.Comments, d.Pos(), fset)
			params := extractParameters(d.Type.Params)
			returnType := extractReturnType(d.Type.Results)

			decl := GoDecl{
				Name:       name,
				Type:       typ,
				Receiver:   recvType,
				FullName:   fullName,
				Parameters: params,
				ReturnType: returnType,
				DocComment: doc,
				IsExported: isExported(name),
				StartLine:  fset.Position(d.Pos()).Line,
				EndLine:    fset.Position(d.End()).Line,
			}

			decls = append(decls, decl)

		case *ast.GenDecl:
			for _, spec := range d.Specs {
				switch s := spec.(type) {
				case *ast.TypeSpec:
					typName := s.Name.Name
					declType := "type"
					var fields []string
					var methods []string

					switch t := s.Type.(type) {
					case *ast.StructType:
						declType = "struct"
						if t.Fields != nil {
							for _, f := range t.Fields.List {
								fieldType := astTypeToString(f.Type)

								if len(f.Names) == 0 {
									// Embedded field
									fields = append(fields, "embedded:"+fieldType)
								} else {
									// Named fields
									for _, n := range f.Names {
										field := n.Name + ": " + fieldType
										if f.Tag != nil {
											field += " " + f.Tag.Value
										}
										fields = append(fields, field)
									}
								}
							}
						}

					case *ast.InterfaceType:
						declType = "interface"
						if t.Methods != nil {
							for _, m := range t.Methods.List {
								if len(m.Names) > 0 {
									// Named method
									for _, name := range m.Names {
										methods = append(methods, name.Name)
									}
								} else {
									// Embedded interface
									embeddedType := astTypeToString(m.Type)
									methods = append(methods, "embedded:"+embeddedType)
								}
							}
						}
					}

					doc := extractDocComment(fileNode.Comments, s.Pos(), fset)
					decl := GoDecl{
						Name:       typName,
						Type:       declType,
						Fields:     fields,
						Methods:    methods,
						DocComment: doc,
						IsExported: isExported(typName),
						StartLine:  fset.Position(d.Pos()).Line,
						EndLine:    fset.Position(d.End()).Line,
					}

					decls = append(decls, decl)

				case *ast.ValueSpec:
					// Variable or constant declarations
					varType := "var"
					if d.Tok == token.CONST {
						varType = "const"
					}

					for _, name := range s.Names {
						doc := extractDocComment(fileNode.Comments, name.Pos(), fset)

						// Get the type if specified
						typeStr := ""
						if s.Type != nil {
							typeStr = astTypeToString(s.Type)
						}

						// Use the individual identifier's position
						startLine := fset.Position(name.Pos()).Line
						endLine := fset.Position(name.End()).Line

						decl := GoDecl{
							Name:       name.Name,
							Type:       varType,
							ReturnType: typeStr, // Reuse return_type field for variable type
							DocComment: doc,
							IsExported: isExported(name.Name),
							StartLine:  startLine,
							EndLine:    endLine,
						}

						decls = append(decls, decl)
					}

				case *ast.ImportSpec:
					// Track imports for graph building
					importPath := ""
					if s.Path != nil {
						importPath = strings.Trim(s.Path.Value, "\"")
					}

					importName := ""
					if s.Name != nil {
						importName = s.Name.Name
					} else {
						// Extract package name from path
						parts := strings.Split(importPath, "/")
						if len(parts) > 0 {
							importName = parts[len(parts)-1]
						}
					}

					if importPath != "" {
						decl := GoDecl{
							Name:      importName,
							Type:      "import",
							FullName:  importPath,
							StartLine: fset.Position(s.Pos()).Line,
							EndLine:   fset.Position(s.End()).Line,
						}
						decls = append(decls, decl)
					}
				}
			}
		}
	}

	// Marshal to JSON and output
	out, err := json.Marshal(decls)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERR: Failed to marshal JSON: %v\n", err)
		os.Exit(4)
	}

	fmt.Println(string(out))
}