package docsite

import (
	"bytes"
	"errors"
	"fmt"
	"go/ast"
	"go/doc"
	"go/parser"
	"go/token"
	"io/fs"
	"os"
	"path"
	"path/filepath"
	"slices"
	"strings"
)

// Ref is a src= reference in a code fence info string: a
// repository-relative path and an optional fragment (a top-level Go
// identifier or a docs:start/docs:end region name).
type Ref struct {
	Path     string
	Fragment string
}

// String returns "path" or "path#fragment".
func (r Ref) String() string {
	if r.Fragment == "" {
		return r.Path
	}
	return r.Path + "#" + r.Fragment
}

// ParseSrc finds the src= field in a fence info string, such as
// "go src=modules/bonfire/example_test.go#ExampleCall".
func ParseSrc(info string) (Ref, bool) {
	for _, f := range strings.Fields(info) {
		if v, ok := strings.CutPrefix(f, "src="); ok {
			p, frag, _ := strings.Cut(v, "#")
			return Ref{Path: p, Fragment: frag}, true
		}
	}
	return Ref{}, false
}

// errSnippetNotFound marks a missing file, identifier or region.
var errSnippetNotFound = errors.New("not found")

// snippetError is a src= reference that breaks a confinement or run rule.
type snippetError string

func (e snippetError) Error() string { return string(e) }

// Extract returns the text a src= reference names, after checking that
// the path is confined to the repository root and that Go sources are code
// go test compiles and runs:
//
//   - no fragment: the whole file;
//   - #region: the lines between "// docs:start region" and
//     "// docs:end region" ("# docs:..." in YAML), dedented;
//   - #Ident: a top-level Go declaration with its doc comment, verbatim;
//     for an Example function, its body dedented with the // Output: line.
func Extract(root string, ref Ref) (string, error) {
	if err := checkRefPath(ref.Path); err != nil {
		return "", err
	}
	rootAbs, err := filepath.Abs(root)
	if err != nil {
		return "", fmt.Errorf("docsite: %w", err)
	}
	rootReal, err := filepath.EvalSymlinks(rootAbs)
	if err != nil {
		return "", fmt.Errorf("docsite: %w", err)
	}
	real, err := filepath.EvalSymlinks(filepath.Join(rootAbs, filepath.FromSlash(ref.Path)))
	if errors.Is(err, fs.ErrNotExist) {
		return "", errSnippetNotFound
	}
	if err != nil {
		return "", fmt.Errorf("docsite: %w", err)
	}
	if !within(real, rootReal) {
		return "", snippetError("path resolves outside the repository root")
	}
	if rel, err := filepath.Rel(rootReal, real); err == nil {
		if err := checkRefPath(filepath.ToSlash(rel)); err != nil {
			return "", snippetError("path resolves to a dotfile or .env file")
		}
	}
	st, err := os.Stat(real)
	if err != nil {
		return "", fmt.Errorf("docsite: %w", err)
	}
	if st.IsDir() {
		return "", snippetError("path is a directory")
	}
	raw, err := os.ReadFile(real)
	if err != nil {
		return "", fmt.Errorf("docsite: %w", err)
	}

	isGo := strings.HasSuffix(real, ".go")
	isTest := strings.HasSuffix(real, "_test.go")
	if isGo {
		if err := checkRootModule(rootReal, filepath.Dir(real)); err != nil {
			return "", err
		}
		if !isTest {
			tests, _ := filepath.Glob(filepath.Join(filepath.Dir(real), "*_test.go"))
			if len(tests) == 0 {
				return "", snippetError("package has no _test.go files, so go test does not cover it")
			}
		}
	}
	if ref.Fragment == "" {
		return string(raw), nil
	}

	lines := strings.Split(string(raw), "\n")
	if text, start, end, found, err := extractRegion(lines, ref.Fragment); err != nil {
		return "", err
	} else if found {
		if isTest {
			if err := checkTestRegion(real, start, end); err != nil {
				return "", err
			}
		}
		return text, nil
	}
	if !isGo {
		return "", errSnippetNotFound
	}
	return extractIdent(real, raw, ref.Fragment, isTest)
}

// checkRefPath enforces the form of a src= path: relative, clean, inside
// the root, and never a dotfile or .env file.
func checkRefPath(p string) error {
	if p == "" {
		return snippetError("src= has no path")
	}
	if filepath.IsAbs(p) || strings.HasPrefix(p, "/") || strings.HasPrefix(p, `\`) || filepath.VolumeName(p) != "" ||
		(len(p) > 1 && p[1] == ':') {
		return snippetError("path must be relative to the repository root")
	}
	if strings.Contains(p, `\`) {
		return snippetError("path must use forward slashes")
	}
	for _, seg := range strings.Split(p, "/") {
		if seg == ".." {
			return snippetError("path must not leave the repository root")
		}
		if strings.HasPrefix(seg, ".") || strings.HasSuffix(seg, ".env") {
			return snippetError("path names a dotfile or .env file")
		}
	}
	if path.Clean(p) != p {
		return snippetError("path must be clean (no ./, // or trailing /)")
	}
	return nil
}

// checkRootModule refuses a Go source that sits in a nested module (a
// directory between it and the root holds a go.mod): root go test does not
// run it.
func checkRootModule(rootReal, dir string) error {
	for d := dir; d != rootReal && within(d, rootReal); d = filepath.Dir(d) {
		if _, err := os.Stat(filepath.Join(d, "go.mod")); err == nil {
			rel, _ := filepath.Rel(rootReal, d)
			return snippetError(fmt.Sprintf("path is inside the nested module %s/go.mod, which root go test does not run", filepath.ToSlash(rel)))
		}
	}
	return nil
}

// extractRegion returns the dedented lines between the docs:start and
// docs:end markers of a region, and the 1-based marker lines.
func extractRegion(lines []string, name string) (text string, start, end int, found bool, err error) {
	isMarker := func(line, kind string) bool {
		t := strings.TrimSpace(line)
		return t == "// docs:"+kind+" "+name || t == "# docs:"+kind+" "+name
	}
	start = -1
	for i, line := range lines {
		if start < 0 && isMarker(line, "start") {
			start = i
			continue
		}
		if start >= 0 && isMarker(line, "end") {
			return strings.Join(dedent(lines[start+1:i]), "\n"), start + 1, i + 1, true, nil
		}
	}
	if start >= 0 {
		return "", 0, 0, false, snippetError(fmt.Sprintf("region %q has no docs:end marker", name))
	}
	return "", 0, 0, false, nil
}

// dedent removes the longest common leading whitespace of the non-blank
// lines and trims leading and trailing blank lines.
func dedent(lines []string) []string {
	for len(lines) > 0 && strings.TrimSpace(lines[0]) == "" {
		lines = lines[1:]
	}
	for len(lines) > 0 && strings.TrimSpace(lines[len(lines)-1]) == "" {
		lines = lines[:len(lines)-1]
	}
	prefix, set := "", false
	for _, l := range lines {
		if strings.TrimSpace(l) == "" {
			continue
		}
		lead := l[:len(l)-len(strings.TrimLeft(l, " \t"))]
		if !set {
			prefix, set = lead, true
			continue
		}
		for !strings.HasPrefix(lead, prefix) {
			prefix = prefix[:len(prefix)-1]
		}
	}
	out := make([]string, len(lines))
	for i, l := range lines {
		if strings.TrimSpace(l) == "" {
			out[i] = ""
			continue
		}
		out[i] = strings.TrimPrefix(l, prefix)
	}
	return out
}

// extractIdent returns a top-level declaration by name, or an Example
// function's body.
func extractIdent(file string, raw []byte, name string, isTest bool) (string, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, file, raw, parser.ParseComments)
	if err != nil {
		return "", snippetError("cannot parse Go source: " + firstLine(err.Error()))
	}
	offset := func(p token.Pos) int { return fset.Position(p).Offset }
	for _, decl := range f.Decls {
		switch d := decl.(type) {
		case *ast.FuncDecl:
			if funcKey(d) != name {
				continue
			}
			if isTest {
				if err := checkTestIdent(file, name, true); err != nil {
					return "", err
				}
			}
			if d.Recv == nil && strings.HasPrefix(name, "Example") && isTest {
				if !exampleHasOutput(f, name) {
					return "", snippetError(name + " has no // Output: comment, so go test compiles it but never runs it")
				}
				body := string(raw[offset(d.Body.Lbrace)+1 : offset(d.Body.Rbrace)])
				return strings.Join(dedent(strings.Split(body, "\n")), "\n"), nil
			}
			start := d.Pos()
			if d.Doc != nil {
				start = d.Doc.Pos()
			}
			return string(raw[offset(start):offset(d.End())]), nil
		case *ast.GenDecl:
			for _, spec := range d.Specs {
				var names []*ast.Ident
				var specDoc *ast.CommentGroup
				switch sp := spec.(type) {
				case *ast.TypeSpec:
					names, specDoc = []*ast.Ident{sp.Name}, sp.Doc
				case *ast.ValueSpec:
					names, specDoc = sp.Names, sp.Doc
				}
				if !slices.ContainsFunc(names, func(id *ast.Ident) bool { return id.Name == name }) {
					continue
				}
				if isTest {
					if err := checkTestIdent(file, name, false); err != nil {
						return "", err
					}
				}
				start, end := d.Pos(), d.End()
				doc := d.Doc
				if d.Lparen.IsValid() {
					start, end, doc = spec.Pos(), spec.End(), specDoc
				}
				if doc != nil {
					start = doc.Pos()
				}
				return string(raw[offset(start):offset(end)]), nil
			}
		}
	}
	return "", errSnippetNotFound
}

// funcKey is "Name" for a function and "Type.Name" for a method.
func funcKey(d *ast.FuncDecl) string {
	if d.Recv == nil || len(d.Recv.List) == 0 {
		return d.Name.Name
	}
	t := d.Recv.List[0].Type
	for {
		switch x := t.(type) {
		case *ast.StarExpr:
			t = x.X
			continue
		case *ast.IndexExpr:
			t = x.X
			continue
		case *ast.IndexListExpr:
			t = x.X
			continue
		case *ast.ParenExpr:
			t = x.X
			continue
		case *ast.Ident:
			return x.Name + "." + d.Name.Name
		}
		return d.Name.Name
	}
}

func exampleHasOutput(f *ast.File, name string) bool {
	for _, ex := range doc.Examples(f) {
		if "Example"+ex.Name == name {
			return ex.Output != "" || ex.EmptyOutput
		}
	}
	return false
}

// testGraph is the reachability of the functions in one directory's test
// files from its Test and Example functions.
type testGraph struct {
	fset       *token.FileSet
	funcs      map[string][]*ast.FuncDecl // by key, per file order
	reachable  map[string]bool            // func keys
	referenced map[string]bool            // identifier names used by reachable funcs
	files      map[string]*ast.File       // by path
}

func isRoot(d *ast.FuncDecl) bool {
	return d.Recv == nil && (strings.HasPrefix(d.Name.Name, "Test") || strings.HasPrefix(d.Name.Name, "Example"))
}

func loadTestGraph(dir string) (*testGraph, error) {
	paths, err := filepath.Glob(filepath.Join(dir, "*_test.go"))
	if err != nil {
		return nil, fmt.Errorf("docsite: %w", err)
	}
	g := &testGraph{fset: token.NewFileSet(), funcs: map[string][]*ast.FuncDecl{},
		reachable: map[string]bool{}, referenced: map[string]bool{}, files: map[string]*ast.File{}}
	byName := map[string][]string{} // plain or method name -> keys
	var queue []string
	for _, p := range paths {
		f, err := parser.ParseFile(g.fset, p, nil, 0)
		if err != nil {
			return nil, snippetError("cannot parse " + filepath.Base(p) + ": " + firstLine(err.Error()))
		}
		g.files[p] = f
		for _, decl := range f.Decls {
			d, ok := decl.(*ast.FuncDecl)
			if !ok {
				continue
			}
			key := funcKey(d)
			g.funcs[key] = append(g.funcs[key], d)
			byName[d.Name.Name] = append(byName[d.Name.Name], key)
			if isRoot(d) && !g.reachable[key] {
				g.reachable[key] = true
				queue = append(queue, key)
			}
		}
	}
	for len(queue) > 0 {
		key := queue[0]
		queue = queue[1:]
		for _, d := range g.funcs[key] {
			if d.Body == nil {
				continue
			}
			ast.Inspect(d.Body, func(n ast.Node) bool {
				id, ok := n.(*ast.Ident)
				if !ok {
					return true
				}
				g.referenced[id.Name] = true
				for _, k := range byName[id.Name] {
					if !g.reachable[k] {
						g.reachable[k] = true
						queue = append(queue, k)
					}
				}
				return true
			})
		}
	}
	return g, nil
}

const notRunMessage = "fragment is not inside a Test or Example function, or a function one of them calls, so go test does not run it"

// checkTestIdent requires a _test.go declaration to be run by go test: a
// function reachable from a Test or Example, or a type, var or const that
// such a function uses.
func checkTestIdent(file, name string, isFunc bool) error {
	g, err := loadTestGraph(filepath.Dir(file))
	if err != nil {
		return err
	}
	if (isFunc && g.reachable[name]) || (!isFunc && g.referenced[name]) {
		return nil
	}
	return snippetError(notRunMessage)
}

// checkTestRegion requires the region's marker lines to sit inside one
// function of the file that a Test or Example function reaches.
func checkTestRegion(file string, start, end int) error {
	g, err := loadTestGraph(filepath.Dir(file))
	if err != nil {
		return err
	}
	f := g.files[file]
	if f == nil {
		return snippetError(notRunMessage)
	}
	for _, decl := range f.Decls {
		d, ok := decl.(*ast.FuncDecl)
		if !ok {
			continue
		}
		from, to := g.fset.Position(d.Pos()).Line, g.fset.Position(d.End()).Line
		if from < start && end < to && g.reachable[funcKey(d)] {
			return nil
		}
	}
	return snippetError(notRunMessage)
}

// drift is a src= fence whose body differs from its source.
type drift struct {
	f    fence
	want string
}

// nestedSrcMessage is the problem for a src= fence that is not a
// top-level block. docs:sync cannot rewrite a line that still carries
// its container marker, so the fence is refused instead of extracted.
const nestedSrcMessage = "src= code block must be a top-level block of the page, not inside a callout, blockquote or list item"

// checkFences verifies every src= fence in fences. lines are the body
// lines the fences index into. lineBase is the 1-based file line of
// lines[0]; file is the display path for problems.
func checkFences(root, file string, lines []string, lineBase int, fences []mdFence) ([]drift, []Problem, error) {
	var drifts []drift
	var problems []Problem
	for _, f := range fences {
		ref, ok := ParseSrc(f.info)
		if !ok {
			continue
		}
		line := lineBase + f.line
		fail := func(msg string) {
			problems = append(problems, Problem{File: file, Line: line, Rule: "snippet", Message: msg})
		}
		if !f.topLevel {
			fail(ref.String() + ": " + nestedSrcMessage)
			continue
		}
		if f.top.close < 0 {
			fail("code block has no closing fence")
			continue
		}
		want, err := Extract(root, ref)
		var se snippetError
		switch {
		case errors.Is(err, errSnippetNotFound):
			fail(ref.String() + " not found")
			continue
		case errors.As(err, &se):
			fail(ref.String() + ": " + se.Error())
			continue
		case err != nil:
			return nil, nil, err
		}
		if fenceBody(lines, f.top) != strings.TrimRight(want, "\n") {
			fail(fmt.Sprintf("body differs from %s (run: summer docs:sync)", ref))
			drifts = append(drifts, drift{f: f.top, want: want})
		}
	}
	return drifts, problems, nil
}

// fenceBody returns a fence's body with the fence indentation removed and
// trailing newlines trimmed.
func fenceBody(lines []string, f fence) string {
	body := make([]string, 0, f.close-f.open)
	for _, l := range lines[f.open+1 : f.close] {
		body = append(body, unindent(l, f.indent))
	}
	return strings.TrimRight(strings.Join(body, "\n"), "\n")
}

func unindent(l string, n int) string {
	i := 0
	for i < n && i < len(l) && l[i] == ' ' {
		i++
	}
	return l[i:]
}

// checkSnippets verifies the src= fences of every docs page. Module
// READMEs are skipped here; a later check refuses src= on those pages
// because their fences are rendered as written.
func (s *site) checkSnippets(docs []parsedDoc) ([]Problem, error) {
	var problems []Problem
	for _, d := range docs {
		if d.page == nil || d.page.Module != "" {
			continue
		}
		fences, err := collectFences(d.doc, d.body)
		if err != nil {
			return nil, err
		}
		_, ps, err := checkFences(s.opts.Root, d.file, strings.Split(string(d.body), "\n"), d.line, fences)
		if err != nil {
			return nil, err
		}
		problems = append(problems, ps...)
	}
	return problems, nil
}

// Sync rewrites every drifted src= fence body under Src from its source,
// preserving everything else byte for byte. When a reference is missing or
// breaks a rule, it returns the problems and writes nothing.
func Sync(opts Options) (SyncResult, []Problem, error) {
	opts, err := opts.normalize()
	if err != nil {
		return SyncResult{}, nil, err
	}
	s := &site{opts: opts}
	files, err := walkPages(opts.Src)
	if err != nil {
		return SyncResult{}, nil, err
	}
	type rewrite struct {
		path string
		mode fs.FileMode
		data []byte
	}
	var rewrites []rewrite
	var problems []Problem
	var result SyncResult
	for _, abs := range files {
		raw, err := os.ReadFile(abs)
		if err != nil {
			return SyncResult{}, nil, fmt.Errorf("docs:sync: read %s: %w", abs, err)
		}
		body, bodyLine := raw, 1
		if _, b, bl, ok := splitFrontmatter(raw); ok {
			body, bodyLine = b, bl
		}
		if !bytes.HasSuffix(raw, body) {
			return SyncResult{}, nil, fmt.Errorf("docs:sync: %s: body is not a suffix of the file", abs)
		}
		doc := s.parseRaw(newMarkdown(), body)
		fences, err := collectFences(doc, body)
		if err != nil {
			return SyncResult{}, nil, err
		}
		lines := strings.Split(string(body), "\n")
		drifts, ps, err := checkFences(opts.Root, s.rel(abs), lines, bodyLine, fences)
		if err != nil {
			return SyncResult{}, nil, err
		}
		for _, p := range ps {
			if !strings.HasPrefix(p.Message, "body differs from ") {
				problems = append(problems, p)
			}
		}
		if len(drifts) == 0 {
			continue
		}
		for i := len(drifts) - 1; i >= 0; i-- {
			d := drifts[i]
			pad := strings.Repeat(" ", d.f.indent)
			var rewritten []string
			for _, l := range strings.Split(strings.TrimRight(d.want, "\n"), "\n") {
				if l == "" {
					rewritten = append(rewritten, "")
				} else {
					rewritten = append(rewritten, pad+l)
				}
			}
			lines = slices.Concat(lines[:d.f.open+1], rewritten, lines[d.f.close:])
		}
		st, err := os.Stat(abs)
		if err != nil {
			return SyncResult{}, nil, fmt.Errorf("docs:sync: %w", err)
		}
		rewritten := append(append([]byte{}, raw[:len(raw)-len(body)]...), []byte(strings.Join(lines, "\n"))...)
		rewrites = append(rewrites, rewrite{path: abs, mode: st.Mode().Perm(), data: rewritten})
		result.Snippets += len(drifts)
		result.Files++
	}
	if len(problems) > 0 {
		sortProblems(problems)
		return SyncResult{}, problems, nil
	}
	for _, r := range rewrites {
		if err := os.WriteFile(r.path, r.data, r.mode); err != nil {
			return SyncResult{}, nil, fmt.Errorf("docs:sync: write %s: %w", r.path, err)
		}
	}
	return result, nil, nil
}