package main

import (
	"fmt"
	"html/template"
	"io/fs"
	"log"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"

	"git.sr.ht/~sircmpwn/getopt"
	"modernc.org/cc/v3"
)

func cutComment(s string) (string, string) {
	start := strings.Index(s, "/*")
	if start < 0 {
		return "", ""
	}

	j := strings.Index(s[start:], "*/")
	if j < 0 {
		return "", ""
	}
	end := start + j + len("*/")

	return s[start:end], s[end:]
}

func extractDocComment(tok cc.Token) []string {
	// Pick the last comment
	rest := tok.Sep.String()
	var comment string
	for {
		nextComment, nextRest := cutComment(rest)
		if nextComment == "" {
			break
		}
		comment, rest = nextComment, nextRest
	}

	if comment == "" || rest != "\n" {
		return nil
	}

	s := comment
	s = strings.TrimPrefix(strings.TrimPrefix(s, "/*"), "*")
	s = strings.TrimSuffix(s, "*/")
	s = strings.TrimSpace(s)

	lines := strings.Split(s, "\n")
	for i, l := range lines {
		l = strings.TrimPrefix(strings.TrimSpace(l), "* ")
		if l == "*" {
			l = ""
		}
		lines[i] = l
	}

	doc := strings.Join(lines, "\n")
	return strings.Split(doc, "\n\n")
}

func enumSpecifierPrototype(enumSpec *cc.EnumSpecifier) prototype {
	var proto prototype
	if enumSpec.Token2.Value != 0 {
		proto.Symbol(declEnum, enumSpec.Token2.Value.String())
	} else {
		proto.Keyword("enum")
	}
	if enumSpec.EnumeratorList != nil {
		proto.Raw(" {\n")
		for el := enumSpec.EnumeratorList; el != nil; el = el.EnumeratorList {
			proto.Raw("\t%v,\n", el.Enumerator.Token)
		}
		proto.Raw("}")
	}
	return proto
}

func structOrUnionSpecifierPrototype(structOrUnionSpec *cc.StructOrUnionSpecifier) prototype {
	kind := structOrUnionSpec.StructOrUnion.Token.Value.String()

	var proto prototype
	if structOrUnionSpec.Token.Value != 0 {
		proto.Symbol(declKind(kind), structOrUnionSpec.Token.Value.String())
	} else {
		proto.Keyword(kind)
	}

	if structOrUnionSpec.StructDeclarationList != nil {
		proto.Raw(" {\n")
		for sdl := structOrUnionSpec.StructDeclarationList; sdl != nil; sdl = sdl.StructDeclarationList {
			proto.Add(indent(structDeclarationPrototype(sdl.StructDeclaration))...).Raw(";\n")
		}
		proto.Raw("}")
	}

	return proto
}

func structDeclarationPrototype(structDecl *cc.StructDeclaration) prototype {
	var proto prototype
	for sql := structDecl.SpecifierQualifierList; sql != nil; sql = sql.SpecifierQualifierList {
		switch sql.Case {
		case cc.SpecifierQualifierListTypeSpec:
			proto.Add(typeSpecifierPrototype(sql.TypeSpecifier)...).Raw(nbsp)
		case cc.SpecifierQualifierListTypeQual:
			proto.Keyword(sql.TypeQualifier.Token.Value.String()).Raw(nbsp)
		default:
			panic(fmt.Sprintf("TODO: %v", sql.Case))
		}
	}
	first := true
	for sdl := structDecl.StructDeclaratorList; sdl != nil; sdl = sdl.StructDeclaratorList {
		if !first {
			proto.Raw(", ")
		}
		proto.Add(declaratorPrototype(sdl.StructDeclarator.Declarator)...)
		first = false
	}
	return proto
}

func declaratorPrototype(declarator *cc.Declarator) prototype {
	var proto prototype
	for ptr := declarator.Pointer; ptr != nil; ptr = ptr.Pointer {
		proto.Raw("*")
		// TODO: ptr.TypeQualifiers
	}
	proto.Add(directDeclaratorPrototype(declarator.DirectDeclarator)...)
	return proto
}

func directDeclaratorPrototype(directDeclarator *cc.DirectDeclarator) prototype {
	var proto prototype
	switch directDeclarator.Case {
	case cc.DirectDeclaratorIdent:
		return *proto.Raw(directDeclarator.Token.Value.String())
	case cc.DirectDeclaratorArr, cc.DirectDeclaratorStaticArr:
		// TODO: TypeQualifiers, AssignmentExpression
		return *proto.Add(directDeclaratorPrototype(directDeclarator.DirectDeclarator)...).Raw("[]")
	case cc.DirectDeclaratorFuncParam:
		ptl := directDeclarator.ParameterTypeList

		proto.Add(directDeclaratorPrototype(directDeclarator.DirectDeclarator)...).Raw("(" + zwsp)
		first := true
		for pl := ptl.ParameterList; pl != nil; pl = pl.ParameterList {
			// TODO: ParameterDeclaration.AttributeSpecifierList
			if !first {
				proto.Raw(", ")
			}
			pd := pl.ParameterDeclaration
			proto.Add(declarationSpecifiersPrototype(pd.DeclarationSpecifiers)...)
			switch pd.Case {
			case cc.ParameterDeclarationDecl:
				proto.Raw(nbsp).Add(declaratorPrototype(pd.Declarator)...)
			case cc.ParameterDeclarationAbstract:
				if pd.AbstractDeclarator != nil {
					panic("TODO: AbstractDeclarator")
				}
			default:
				panic(fmt.Sprintf("unsupported %v", pd.Case))
			}
			first = false
		}
		if ptl.Case == cc.ParameterTypeListVar {
			proto.Raw(", ...")
		}
		proto.Raw(")")
		return proto
	case cc.DirectDeclaratorDecl:
		// TODO: AttributeSpecifierList
		return *proto.Raw("(").Add(declaratorPrototype(directDeclarator.Declarator)...).Raw(")")
	default:
		panic(fmt.Sprintf("TODO: %v", directDeclarator.Case))
	}
}

func declarationSpecifiersPrototype(declSpec *cc.DeclarationSpecifiers) prototype {
	var proto prototype
	first := true
	for declSpec := declSpec; declSpec != nil; declSpec = declSpec.DeclarationSpecifiers {
		if !first {
			proto.Raw(nbsp)
		}
		switch declSpec.Case {
		case cc.DeclarationSpecifiersTypeSpec:
			proto.Add(typeSpecifierPrototype(declSpec.TypeSpecifier)...)
		case cc.DeclarationSpecifiersTypeQual:
			proto.Keyword(declSpec.TypeQualifier.Token.Value.String())
		default:
			panic(fmt.Sprintf("TODO: %v", declSpec.Case))
		}
		first = false
	}
	return proto
}

func typeSpecifierPrototype(typeSpec *cc.TypeSpecifier) prototype {
	var proto prototype
	switch typeSpec.Case {
	case cc.TypeSpecifierStructOrUnion:
		return structOrUnionSpecifierPrototype(typeSpec.StructOrUnionSpecifier)
	case cc.TypeSpecifierEnum:
		return enumSpecifierPrototype(typeSpec.EnumSpecifier)
	case cc.TypeSpecifierAtomic:
		panic("TODO")
	case cc.TypeSpecifierBool:
		return *proto.Keyword("bool")
	default:
		return *proto.Keyword(typeSpec.Token.Value.String())
	}
}

func typeSpecifierFirstToken(typeSpec *cc.TypeSpecifier) cc.Token {
	switch typeSpec.Case {
	case cc.TypeSpecifierStructOrUnion:
		return typeSpec.StructOrUnionSpecifier.StructOrUnion.Token
	case cc.TypeSpecifierEnum:
		return typeSpec.EnumSpecifier.Token
	case cc.TypeSpecifierAtomic:
		return typeSpec.AtomicTypeSpecifier.Token
	default:
		return typeSpec.Token
	}
}

func indent(in prototype) prototype {
	var out prototype
	out.Raw("\t")
	for _, tok := range in {
		if tok.Type == tokenRaw {
			tok.Value = strings.ReplaceAll(tok.Value, "\n", "\n\t")
		}
		out = append(out, tok)
	}
	return out
}

func isExportedSymbol(l []string, sym string) bool {
	if len(l) == 0 {
		return true
	}

	for _, pattern := range l {
		ok, err := path.Match(pattern, sym)
		if err != nil {
			log.Fatalf("failed to match exported symbol pattern %q: %v", pattern, err)
		} else if ok {
			return true
		}
	}

	return false
}

func applyFilePrefixMap(filename string, m [][2]string) string {
	for _, entry := range m {
		old, new := entry[0], entry[1]
		if strings.HasPrefix(filename, old) {
			return strings.Replace(filename, old, new, 1)
		}
	}
	return filename
}

func relPath(base, target string) string {
	baseElems := strings.Split(path.Clean(base), "/")
	targetElems := strings.Split(path.Clean(target), "/")

	// Skip common elements
	i := 0
	for i < len(baseElems) && i < len(targetElems) && baseElems[i] == targetElems[i] {
		i++
	}

	var out []string
	for j := i; j < len(baseElems)-1; j++ {
		out = append(out, "..")
	}
	for j := i; j < len(targetElems); j++ {
		out = append(out, targetElems[j])
	}
	return path.Join(out...)
}

// TODO: replace with strings.Cut once Go 1.18 is widespread
func cut(s, sep string) (before, after string, found bool) {
	split := strings.SplitN(s, sep, 2)
	if len(split) == 2 {
		return split[0], split[1], true
	} else {
		return s, "", false
	}
}

func renderTemplate(tpl *template.Template, outputFilename, name string, data interface{}) {
	if err := os.MkdirAll(filepath.Dir(outputFilename), 0755); err != nil {
		log.Fatalf("failed to create output directory: %v", err)
	}

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

	if err := tpl.ExecuteTemplate(f, name, data); err != nil {
		log.Fatal(err)
	}
}

func main() {
	opts, optind, err := getopt.Getopts(os.Args, "I:D:U:f:o:")
	if err != nil {
		log.Fatal(err)
	}
	targets := os.Args[optind:]

	if len(targets) == 0 {
		log.Fatal("no input files specified")
	}

	var D, U, I, exportedSymbols []string
	outputDir := "out"
	var filePrefixMap [][2]string
	for _, opt := range opts {
		switch opt.Option {
		case 'I':
			I = append(I, opt.Value)
		case 'D':
			D = append(D, opt.Value)
		case 'U':
			U = append(U, opt.Value)
		case 'f':
			name, value, _ := cut(opt.Value, "=")
			switch name {
			case "exported-symbols":
				exportedSymbols = append(exportedSymbols, value)
			case "file-prefix-map":
				old, new, ok := cut(value, "=")
				if !ok {
					log.Fatal("invalid syntax in -ffile-prefix-map")
				}
				filePrefixMap = append(filePrefixMap, [2]string{old, new})
			default:
				log.Fatalf("unknown -f option: %v", name)
			}
		case 'o':
			outputDir = opt.Value
		}
	}

	hostPredefined, hostIncludePaths, hostSysIncludePaths, err := cc.HostConfig("")
	if err != nil {
		log.Fatalf("failed to get host config: %v", err)
	}

	sources := []cc.Source{
		{Name: "<predefined>", Value: hostPredefined},
		{Name: "<builtin>", Value: builtin},
	}

	if s := generateDefines(D); s != "" {
		sources = append(sources, cc.Source{Name: "<defines>", Value: s})
	}
	if s := generateUndefines(U); s != "" {
		sources = append(sources, cc.Source{Name: "<undefines>", Value: s})
	}

	targetMap := make(map[string]string, len(targets))
	for _, target := range targets {
		fi, err := os.Stat(target)
		if err != nil {
			log.Fatalf("failed to stat input file %q: %v", target, err)
		}

		var filenames []string
		if fi.IsDir() {
			err := filepath.WalkDir(target, func(path string, de fs.DirEntry, err error) error {
				if err != nil {
					return err
				}
				if !de.Type().IsRegular() || filepath.Ext(path) != ".h" {
					return nil
				}
				filenames = append(filenames, path)
				return nil
			})
			if err != nil {
				log.Fatalf("failed to walk dir %q: %v", target, err)
			}
		} else {
			filenames = []string{target}
		}

		for _, filename := range filenames {
			origFilename := applyFilePrefixMap(filename, filePrefixMap)

			// cc evaluates symlinks
			filename, err := filepath.EvalSymlinks(filename)
			if err != nil {
				log.Fatalf("failed to eval symlinks: %v", err)
			}

			// cc expects absolute paths
			filename, err = filepath.Abs(filename)
			if err != nil {
				log.Fatalf("failed to get absolute path: %v", err)
			}

			targetMap[filename] = origFilename
			sources = append(sources, cc.Source{Name: filename})
		}
	}

	// Priority order is relative (for #include "..." only), then -I, then the
	// usual places. See:
	// https://pubs.opengroup.org/onlinepubs/9699919799/utilities/c99.html
	includePaths := append([]string{"@"}, I...)
	includePaths = append(includePaths, hostIncludePaths...)
	includePaths = append(includePaths, hostSysIncludePaths...)

	sysIncludePaths := append([]string(nil), I...)
	sysIncludePaths = append(sysIncludePaths, hostSysIncludePaths...)

	cfg := new(cc.Config)
	cfg.PreserveWhiteSpace = true
	ast, err := cc.Parse(cfg, includePaths, sysIncludePaths, sources)
	if err != nil {
		log.Fatalf("parse error: %v", err)
	}

	var decls []declData
	for tu := ast.TranslationUnit; tu != nil; tu = tu.TranslationUnit {
		if tu.ExternalDeclaration == nil {
			continue
		}
		if tu.ExternalDeclaration.Case != cc.ExternalDeclarationDecl {
			continue
		}
		decl := tu.ExternalDeclaration.Declaration

		origFilename, ok := targetMap[decl.Position().Filename]
		if !ok {
			continue
		}

		if decl.DeclarationSpecifiers.Case != cc.DeclarationSpecifiersTypeSpec {
			continue
		}
		typeSpec := decl.DeclarationSpecifiers.TypeSpecifier
		typeSpecProto := typeSpecifierPrototype(typeSpec)

		if decl.InitDeclaratorList == nil {
			proto := append(typeSpecProto, token{Value: ";"})
			switch typeSpec.Case {
			case cc.TypeSpecifierStructOrUnion:
				structOrUnionSpec := typeSpec.StructOrUnionSpecifier
				kind := structOrUnionSpec.StructOrUnion.Token.Value.String()
				decls = append(decls, declData{
					Kind:        declKind(kind),
					Name:        structOrUnionSpec.Token.Value.String(),
					Prototype:   proto,
					Description: extractDocComment(structOrUnionSpec.StructOrUnion.Token),
					Filename:    origFilename,
				})
			case cc.TypeSpecifierEnum:
				enumSpec := typeSpec.EnumSpecifier
				decls = append(decls, declData{
					Kind:        declEnum,
					Name:        enumSpec.Token2.Value.String(),
					Prototype:   proto,
					Description: extractDocComment(enumSpec.Token),
					Filename:    origFilename,
				})
			}
		} else {
			for idl := decl.InitDeclaratorList; idl != nil; idl = idl.InitDeclaratorList {
				var kind declKind
				switch idl.InitDeclarator.Declarator.DirectDeclarator.Case {
				case cc.DirectDeclaratorIdent:
					kind = declVar
				case cc.DirectDeclaratorFuncParam:
					kind = declFunc
				default:
					panic(fmt.Sprintf("TODO: %v", idl.InitDeclarator.Declarator.DirectDeclarator.Case))
				}
				var proto prototype
				proto.Add(typeSpecProto...).Raw(" ").Add(declaratorPrototype(idl.InitDeclarator.Declarator)...).Raw(";")
				decls = append(decls, declData{
					Kind:        kind,
					Name:        idl.InitDeclarator.Declarator.Name().String(),
					Prototype:   proto,
					Description: extractDocComment(typeSpecifierFirstToken(typeSpec)),
					Filename:    origFilename,
				})
			}
		}
	}

	declMap := make(map[declKey]declData)
	for _, decl := range decls {
		if prev, ok := declMap[decl.key()]; ok {
			// Prefer the declaration with a longer prototype. In case of a
			// tie (e.g. because both are opaque structs), prefer the
			// declaration with a longer description.
			if len(decl.Prototype) < len(prev.Prototype) {
				continue
			} else if decl.Prototype == nil && len(decl.Description) < len(prev.Description) {
				continue
			}
		}
		if !isExportedSymbol(exportedSymbols, decl.Name) {
			continue
		}

		declMap[decl.key()] = decl
	}

	declByFile := make(map[string][]declData)
	for _, decl := range declMap {
		declByFile[decl.Filename] = append(declByFile[decl.Filename], decl)
	}

	tpl := template.New("gyosu")
	tpl.Funcs(template.FuncMap{
		"symhref": func(kind declKind, name string) string {
			panic("Unimplemented")
		},
		"unithref": func(name string) string {
			return filepath.ToSlash(name) + ".html"
		},
	})
	template.Must(tpl.ParseGlob("template/*.html"))

	if err := os.RemoveAll(outputDir); err != nil {
		log.Fatalf("failed to remove %q: %v", outputDir, err)
	}

	data := indexData{Title: "Documentation"}
	for filename, _ := range declByFile {
		data.Filenames = append(data.Filenames, filename)
	}
	sort.Strings(data.Filenames)

	indexFilename := filepath.Join(outputDir, "index.html")
	renderTemplate(tpl, indexFilename, "index.html", &data)

	for filename, decls := range declByFile {
		filename := filename // capture variable

		basePath := filepath.ToSlash(filename + ".html")

		data := unitData{
			Title:     "Documentation for <" + filename + ">",
			Decls:     decls,
			IndexHref: relPath(basePath, ""),
		}

		sort.Slice(data.Decls, func(i, j int) bool {
			return data.Decls[i].Name < data.Decls[j].Name
		})

		symHref := func(kind declKind, name string) string {
			k := declKey{Kind: kind, Name: name}
			decl, ok := declMap[k]
			if !ok {
				return ""
			}

			anchor := fmt.Sprintf("#%v-%v", kind, name)
			if decl.Filename == filename {
				return anchor
			}

			targetPath := filepath.ToSlash(decl.Filename + ".html")
			return relPath(basePath, targetPath) + anchor
		}

		unitHref := func(name string) string {
			panic("Unimplemented")
		}

		tpl.Funcs(template.FuncMap{
			"symhref":  symHref,
			"unithref": unitHref,
		})

		outputFilename := filepath.Join(outputDir, filename+".html")
		renderTemplate(tpl, outputFilename, "unit.html", &data)
	}
}