// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"maps"
	"os"
	"path/filepath"
	"reflect"
	"slices"
	"strings"
	"text/template"

	"github.com/siderolabs/gen/xslices"

	"gopkg.in/yaml.v3"
	"mvdan.cc/gofumpt/format"
)

var tpl = `// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Code generated by hack/docgen tool. DO NOT EDIT.

package {{ .Package }}

import (
	"github.com/siderolabs/talos/pkg/machinery/config/encoder"
)

{{ $tick := "` + "`" + `" -}}

{{ range $struct := .Structs -}}
func ({{ $struct.Name }}) Doc() *encoder.Doc {
	doc := &encoder.Doc{
		Type : "{{ if $struct.Text.Alias }}{{ $struct.Text.Alias}}{{ else }}{{ $struct.Name }}{{ end }}",
		Comments: [3]string{ "" /* encoder.HeadComment */, "{{ $struct.Text.Comment }}" /* encoder.LineComment */,  "" /* encoder.FootComment */},
		Description: "{{ $struct.Text.Description }}",
		{{ if $struct.AppearsIn -}}
		AppearsIn: []encoder.Appearance{
		{{ range $value := $struct.AppearsIn -}}
			{
				TypeName: "{{ $value.Struct.Name }}",
				FieldName: "{{ $value.FieldName }}",
			},
		{{ end -}}
		},
		{{ end -}}
		Fields: []encoder.Doc{
			{{ range $index, $field := $struct.Fields -}}
				{{ if $field.Tag -}}
					{
					  Name: "{{ $field.Tag }}",
					  Type: "{{ $field.Type }}",
					  Note: "{{ $field.Note }}",
					  Description: "{{ $field.Text.Description }}",
					  Comments: [3]string{ "" /* encoder.HeadComment */, "{{ $field.Text.Comment }}" /* encoder.LineComment */,  "" /* encoder.FootComment */},
	  				  {{ if $field.Text.Values -}}
					  Values : []string{
						{{ range $value := $field.Text.Values -}}
							"{{ $value }}",
						{{ end -}}
					  },
					  {{ end -}}
					},
				{{ else -}}
					{},
				{{- end }}
			{{- end }}

		},
	}

	{{ range $example := $struct.Text.Examples }}
	{{ if $example.Value }}
	doc.AddExample("{{ $example.Name }}", {{ $example.Value }})
	{{ end -}}
	{{ end }}

	{{ range $index, $field := $struct.Fields -}}
		{{ if $field.Tag -}}
			{{ if $field.Text.Examples -}}
				{{ range $example := $field.Text.Examples -}}
					{{- if $example.Value }}
						doc.Fields[{{ $index }}].AddExample("{{ $example.Name }}", {{ $example.Value }})
					{{- end }}
				{{- end }}
			{{- end }}
		{{- end }}
	{{- end }}


	return doc
}
{{ end -}}

// GetFileDoc returns documentation for the file {{ .File }}.
func GetFileDoc() *encoder.FileDoc {
	return &encoder.FileDoc{
		Name: "{{ .Name }}",
		Description: "{{ .Header }}",
		Structs: []*encoder.Doc{
			{{ range $struct := .Structs -}}
			{{ $struct.Name }}{}.Doc(),
			{{ end -}}
		},
	}
}
`

type Doc struct {
	Name    string
	Package string
	Title   string
	Header  string
	File    string
	Structs []*Struct
}

type Struct struct {
	Name      string
	Text      *Text
	Fields    []*Field
	AppearsIn []Appearance
}

type Appearance struct {
	Struct    *Struct
	FieldName string
}

type Example struct {
	Name  string `yaml:"name"`
	Value string `yaml:"value"`
}

type Field struct {
	Name    string
	Type    string
	TypeRef string
	Text    *Text
	Tag     string
	Note    string
}

type Text struct {
	Comment        string         `json:"-"`
	Description    string         `json:"description"`
	Examples       []*Example     `json:"examples"`
	Alias          string         `json:"alias"`
	Values         []string       `json:"values"`
	Schema         *SchemaWrapper `json:"schema"`
	SchemaRoot     bool           `json:"schemaRoot" yaml:"schemaRoot"`
	SchemaRequired bool           `json:"schemaRequired" yaml:"schemaRequired"`
	SchemaMeta     string         `json:"schemaMeta" yaml:"schemaMeta"`
}

func in(p string) (string, error) {
	return filepath.Abs(p)
}

func out(p string) (*os.File, error) {
	abs, err := filepath.Abs(p)
	if err != nil {
		return nil, err
	}

	return os.Create(abs)
}

type packageType struct {
	name    string
	doc     string
	file    string
	structs []*structType
}

type structType struct {
	name string
	text *Text
	pos  token.Pos
	node *ast.StructType
}

type aliasType struct {
	fieldType    string
	fieldTypeRef string
}

func collectStructs(node ast.Node) ([]*structType, map[string]aliasType) {
	var structs []*structType

	aliases := map[string]aliasType{}

	collectStructs := func(n ast.Node) bool {
		g, ok := n.(*ast.GenDecl)
		if !ok {
			return true
		}

		isAlias := false

		if g.Doc != nil {
			for _, comment := range g.Doc.List {
				if strings.Contains(comment.Text, "docgen:nodoc") {
					return true
				}

				if strings.Contains(comment.Text, "docgen:alias") {
					isAlias = true
				}
			}
		}

		for _, spec := range g.Specs {
			t, ok := spec.(*ast.TypeSpec)
			if !ok {
				return true
			}

			if t.Type == nil {
				return true
			}

			x, ok := t.Type.(*ast.StructType)
			if !ok {
				if isAlias {
					aliases[t.Name.Name] = aliasType{
						fieldType:    formatFieldType(t.Type),
						fieldTypeRef: getFieldType(t.Type),
					}
				}

				return true
			}

			structName := t.Name.Name

			text := &Text{}

			if t.Doc != nil {
				text = parseComment([]byte(t.Doc.Text()))
			} else if g.Doc != nil {
				text = parseComment([]byte(g.Doc.Text()))
			}

			s := &structType{
				name: structName,
				text: text,
				node: x,
				pos:  x.Pos(),
			}

			structs = append(structs, s)
		}

		return true
	}

	ast.Inspect(node, collectStructs)

	return structs, aliases
}

func parseComment(comment []byte) *Text {
	text := &Text{}
	if err := yaml.Unmarshal(comment, text); err != nil {
		lines := strings.Split(string(comment), "\n")
		for i := range lines {
			lines[i] = strings.TrimLeft(lines[i], "\t")
		}

		// not yaml, fallback
		text.Description = strings.Join(lines, "\n")
		// take only the first line from the Description for the comment
		text.Comment = lines[0]

		// try to parse everything except for the first line as yaml
		if err = yaml.Unmarshal([]byte(strings.Join(lines[1:], "\n")), text); err == nil {
			// if parsed, remove it from the description
			text.Description = text.Comment
		}
	} else {
		text.Description = strings.TrimSpace(text.Description)
		// take only the first line from the Description for the comment
		text.Comment = strings.Split(text.Description, "\n")[0]
	}

	text.Comment = escape(text.Comment)

	text.Description = escape(text.Description)
	for _, example := range text.Examples {
		example.Name = escape(example.Name)
		example.Value = strings.TrimSpace(example.Value)
	}

	return text
}

func getFieldType(p any) string {
	if m, ok := p.(*ast.MapType); ok {
		return getFieldType(m.Value)
	}

	switch t := p.(type) {
	case *ast.Ident:
		return t.Name
	case *ast.ArrayType:
		return getFieldType(p.(*ast.ArrayType).Elt)
	case *ast.StarExpr:
		return getFieldType(t.X)
	case *ast.SelectorExpr:
		return getFieldType(t.Sel)
	default:
		return ""
	}
}

func formatFieldType(p any) string {
	if m, ok := p.(*ast.MapType); ok {
		return fmt.Sprintf("map[%s]%s", formatFieldType(m.Key), formatFieldType(m.Value))
	}

	switch t := p.(type) {
	case *ast.Ident:
		return t.Name
	case *ast.ArrayType:
		return "[]" + formatFieldType(p.(*ast.ArrayType).Elt)
	case *ast.StructType:
		return "struct"
	case *ast.StarExpr:
		return formatFieldType(t.X)
	case *ast.SelectorExpr:
		return formatFieldType(t.Sel)
	default:
		log.Printf("unknown: %#v", t)

		return ""
	}
}

func escape(value string) string {
	value = strings.ReplaceAll(value, `"`, `\"`)
	value = strings.ReplaceAll(value, "\n", `\n`)

	return strings.TrimSpace(value)
}

func collectFields(s *structType, aliases map[string]aliasType) (fields []*Field) {
	fields = []*Field{}

	for _, f := range s.node.Fields.List {
		if f.Names == nil {
			// This is an embedded struct.
			fields = append(fields, &Field{Type: "unknown"})

			continue
		}

		name := f.Names[0].Name

		if f.Doc == nil {
			log.Fatalf("field %q is missing a documentation", name)
		}

		if strings.Contains(f.Doc.Text(), "docgen:nodoc") {
			fields = append(fields, &Field{Type: "unknown"})

			continue
		}

		if f.Tag == nil {
			log.Fatalf("field %q is missing a tag", name)
		}

		fieldType := formatFieldType(f.Type)

		if alias, ok := aliases[fieldType]; ok {
			fieldType = alias.fieldType
		}

		fieldTypeRef := getFieldType(f.Type)

		if alias, ok := aliases[fieldTypeRef]; ok {
			fieldTypeRef = alias.fieldTypeRef
		}

		tag := reflect.StructTag(strings.Trim(f.Tag.Value, "`"))
		yamlTag := tag.Get("yaml")
		yamlTag = strings.Split(yamlTag, ",")[0]

		if yamlTag == "" {
			log.Fatalf("field %q is missing a `yaml` tag or name in it", name)
		}

		text := parseComment([]byte(f.Doc.Text()))

		field := &Field{
			Name:    name,
			Tag:     yamlTag,
			Type:    fieldType,
			TypeRef: fieldTypeRef,
			Text:    text,
		}

		if f.Comment != nil {
			field.Note = escape(f.Comment.Text())
		}

		fields = append(fields, field)
	}

	return fields
}

func renderDoc(doc *Doc, dest string) {
	t := template.Must(template.New("docfile.tpl").Parse(tpl))
	buf := bytes.Buffer{}

	err := t.Execute(&buf, doc)
	if err != nil {
		log.Fatalf("failed to render template: %v", err)
	}

	formatted, err := format.Source(buf.Bytes(), format.Options{})
	if err != nil {
		log.Printf("data: %s", buf.Bytes())
		log.Fatalf("failed to format source: %v", err)
	}

	out, err := out(dest)
	if err != nil {
		log.Fatalf("failed to create output file: %v", err)
	}
	defer out.Close()

	_, err = out.Write(formatted)
	if err != nil {
		log.Fatalf("failed to write output file: %v", err)
	}
}

func processFiles(inputFiles []string, outputFile, schemaOutputFile, versionTagFile string) {
	var packageNames []string

	packageNameToType := map[string]*packageType{}
	aliases := map[string]aliasType{}

	for _, inputFile := range inputFiles {
		abs, err := in(inputFile)
		if err != nil {
			log.Fatal(err)
		}

		log.Printf("creating package file set: %q", abs)

		fset := token.NewFileSet()

		node, err := parser.ParseFile(fset, abs, nil, parser.ParseComments)
		if err != nil {
			log.Fatal(err)
		}

		packageName := node.Name.Name

		if _, ok := packageNameToType[packageName]; !ok {
			packageNameToType[packageName] = &packageType{
				name: packageName,
				file: outputFile,
			}

			packageNames = append(packageNames, packageName)
		}

		pkg := packageNameToType[packageName]

		if node.Doc != nil && node.Doc.Text() != "" {
			pkg.doc = node.Doc.Text()
		}

		tokenFile := fset.File(node.Pos())
		if tokenFile == nil {
			log.Fatalf("No token")
		}

		log.Printf("parsing file in package %q: %s", packageName, tokenFile.Name())

		fileStructs, fileAliases := collectStructs(node)

		pkg.structs = append(pkg.structs, fileStructs...)

		maps.Copy(aliases, fileAliases)
	}

	slices.Sort(packageNames)

	docs := xslices.Map(packageNames, func(name string) *Doc {
		return packageToDoc(packageNameToType[name], aliases)
	})

	if schemaOutputFile != "" {
		renderSchema(docs, schemaOutputFile, versionTagFile)
	}

	if outputFile == "" {
		return
	}

	if len(docs) != 1 {
		log.Fatalf("expected exactly one package to generate docs, got %d", len(docs))
	}

	renderDoc(docs[0], outputFile)
}

func packageToDoc(pkg *packageType, aliases map[string]aliasType) *Doc {
	if len(pkg.structs) == 0 {
		log.Fatalf("failed to find types that could be documented in %v", pkg.file)
	}

	doc := &Doc{
		Package: pkg.name,
		Structs: []*Struct{},
	}

	extraExamples := map[string][]*Example{}
	backReferences := map[string][]Appearance{}

	for _, s := range pkg.structs {
		log.Printf("generating docs for type: %q", s.name)

		fields := collectFields(s, aliases)

		s := &Struct{
			Name:   s.name,
			Text:   s.text,
			Fields: fields,
		}

		for _, field := range fields {
			if field.TypeRef == "" {
				continue
			}

			if len(field.Text.Examples) > 0 {
				extraExamples[field.TypeRef] = append(extraExamples[field.TypeRef], field.Text.Examples...)
			}

			backReferences[field.TypeRef] = append(backReferences[field.TypeRef], Appearance{
				Struct:    s,
				FieldName: field.Tag,
			})
		}

		doc.Structs = append(doc.Structs, s)
	}

	for _, s := range doc.Structs {
		if s.Text.Alias != "" {
			s.Text.Description = strings.ReplaceAll(s.Text.Description, s.Name, s.Text.Alias)
			s.Text.Comment = strings.ReplaceAll(s.Text.Comment, s.Name, s.Text.Alias)
		}

		if extra, ok := extraExamples[s.Name]; ok {
			s.Text.Examples = append(s.Text.Examples, extra...)
		}

		if ref, ok := backReferences[s.Name]; ok {
			s.AppearsIn = append(s.AppearsIn, ref...)
		}
	}

	doc.Package = pkg.name
	doc.Name = doc.Package
	doc.Header = escape(pkg.doc)

	doc.File = pkg.file

	return doc
}

func sourcesWithJSONSchema(dir string) []string {
	var sources []string

	if err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() || !strings.HasSuffix(info.Name(), ".go") {
			return nil
		}

		fileBytes, err := os.ReadFile(path)
		if err != nil {
			return err
		}

		if strings.Contains(string(fileBytes), "//docgen:jsonschema") {
			sources = append(sources, path)
		}

		return nil
	}); err != nil {
		log.Fatalf("failed to walk directory %q: %v", dir, err)
	}

	return sources
}

func determineFiles(generateSchemaFromDir string, args []string) []string {
	if generateSchemaFromDir != "" {
		if len(args) > 0 {
			log.Fatalf("cannot specify both -generate-schema-from-dir and input files as args")
		}

		files := sourcesWithJSONSchema(generateSchemaFromDir)

		if len(files) == 0 {
			log.Fatalf("no Go files annotated with //docgen:jsonschema found in %q", generateSchemaFromDir)
		}

		return files
	}

	if len(args) == 0 {
		log.Fatalf("no input files")
	}

	return args
}

func main() {
	outputFile := flag.String("output", "", "output file name")
	jsonSchemaOutputFile := flag.String("json-schema-output", "", "output file name for json schema")
	versionTagFile := flag.String("version-tag-file", "", "file name for version tag")
	generateSchemaFromDir := flag.String("generate-schema-from-dir", "", "generate a JSON schema by recursively parsing the sources in the specified directory")

	flag.Parse()

	files := determineFiles(*generateSchemaFromDir, flag.Args())

	processFiles(files, *outputFile, *jsonSchemaOutputFile, *versionTagFile)
}