package main

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"html/template"
	"io"
	"io/fs"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"slices"
	"sort"
	"strings"
	"time"

	"github.com/yuin/goldmark"
	"github.com/yuin/goldmark/ast"
	"github.com/yuin/goldmark/extension"
	"github.com/yuin/goldmark/parser"
	gmtext "github.com/yuin/goldmark/text"
	"golang.org/x/net/html"
	"gopkg.in/yaml.v3"
)

const manifestName = ".weft-generated.json"

type Page struct {
	SourcePath   string
	OutputPath   string
	URL          string
	CanonicalURL string
	Title        string
	Meta         map[string]any
	ModTime      time.Time
	Sections     []template.HTML
}

type templateData struct {
	*Page
	Pages map[string][]*Page
}

type source struct {
	path   string
	output string
	kind   string
	page   *Page
}

type site struct {
	root      string
	pages     map[string][]*Page
	sources   []source
	templates *template.Template
	warnings  []string
}

type manifest struct {
	Outputs []string `json:"outputs"`
}

func main() {
	if len(os.Args) != 2 {
		fmt.Fprintln(os.Stderr, "usage: weft <site-root>")
		os.Exit(2)
	}
	if err := build(os.Args[1]); err != nil {
		fmt.Fprintln(os.Stderr, "weft:", err)
		os.Exit(1)
	}
}

func build(root string) error {
	abs, err := filepath.Abs(root)
	if err != nil {
		return err
	}
	st, err := os.Stat(abs)
	if err != nil {
		return err
	}
	if !st.IsDir() {
		return fmt.Errorf("site root is not a directory: %s", abs)
	}

	s := &site{root: abs, pages: map[string][]*Page{}}
	if err := s.discover(); err != nil {
		return err
	}
	if err := s.loadTemplates(); err != nil {
		return err
	}
	outputs, err := s.render()
	if err != nil {
		return err
	}
	if err := s.validate(outputs); err != nil {
		return err
	}
	if err := s.write(outputs); err != nil {
		return err
	}
	for _, warning := range s.warnings {
		fmt.Fprintln(os.Stderr, "weft: warning:", warning)
	}
	return nil
}

func excluded(rel string) bool {
	for _, part := range strings.Split(filepath.ToSlash(rel), "/") {
		if strings.HasPrefix(part, "_") || part == ".git" || part == ".jj" || strings.HasPrefix(part, ".weft-txn-") {
			return true
		}
	}
	return false
}

func (s *site) discover() error {
	outputs := map[string]string{}
	var templatePaths []string
	err := filepath.WalkDir(s.root, func(full string, entry fs.DirEntry, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}
		rel, err := filepath.Rel(s.root, full)
		if err != nil || rel == "." {
			return err
		}
		if excluded(rel) {
			if entry.IsDir() {
				return filepath.SkipDir
			}
			return nil
		}
		if entry.IsDir() {
			return nil
		}
		rel = filepath.ToSlash(rel)
		info, err := entry.Info()
		if err != nil {
			return err
		}
		switch {
		case strings.HasSuffix(rel, ".md"):
			out := strings.TrimSuffix(rel, ".md") + ".html"
			if err := claim(outputs, out, rel); err != nil {
				return err
			}
			page, warnings, err := s.parsePage(rel, out, info)
			if err != nil {
				return fmt.Errorf("parse %s: %w", rel, err)
			}
			s.warnings = append(s.warnings, warnings...)
			s.sources = append(s.sources, source{path: rel, output: out, kind: "markdown", page: page})
			s.addPage(page)
		case strings.HasSuffix(rel, ".tmpl"):
			templatePaths = append(templatePaths, rel)
			var out string
			if strings.HasSuffix(rel, ".html.tmpl") || strings.HasSuffix(rel, ".xml.tmpl") {
				out = strings.TrimSuffix(rel, ".tmpl")
				if err := claim(outputs, out, rel); err != nil {
					return err
				}
				s.sources = append(s.sources, source{path: rel, output: out, kind: "template"})
			}
		}
		return nil
	})
	if err != nil {
		return err
	}
	slices.Sort(templatePaths)
	// Keep all template paths as non-output sources so definitions are loaded.
	for _, p := range templatePaths {
		found := false
		for _, src := range s.sources {
			if src.path == p {
				found = true
				break
			}
		}
		if !found {
			s.sources = append(s.sources, source{path: p, kind: "partial"})
		}
	}
	sort.SliceStable(s.sources, func(i, j int) bool { return s.sources[i].path < s.sources[j].path })
	return nil
}

func claim(outputs map[string]string, output, input string) error {
	if previous, ok := outputs[output]; ok {
		return fmt.Errorf("output collision: %s and %s both target %s", previous, input, output)
	}
	outputs[output] = input
	return nil
}

func (s *site) parsePage(rel, out string, info fs.FileInfo) (*Page, []string, error) {
	raw, err := os.ReadFile(filepath.Join(s.root, filepath.FromSlash(rel)))
	if err != nil {
		return nil, nil, err
	}
	meta, body, err := frontmatter(raw)
	if err != nil {
		return nil, nil, err
	}
	md := goldmark.New(goldmark.WithExtensions(extension.GFM), goldmark.WithParserOptions(parser.WithAutoHeadingID()))
	doc := md.Parser().Parse(gmtext.NewReader(body))
	rewriteLinks(doc)
	title := extractTitle(doc, body)
	var warnings []string
	if title == "" {
		if value, ok := meta["title"].(string); ok && strings.TrimSpace(value) != "" {
			title = strings.TrimSpace(value)
		} else {
			base := strings.TrimSuffix(path.Base(rel), path.Ext(rel))
			title = strings.Title(strings.ReplaceAll(strings.ReplaceAll(base, "_", " "), "-", " ")) //nolint:staticcheck
			warnings = append(warnings, fmt.Sprintf("%s has no H1; using filename-derived title %q", rel, title))
		}
	}
	sections, err := renderSections(md, doc, body)
	if err != nil {
		return nil, nil, err
	}
	urlPath := "/" + out
	if path.Base(out) == "index.html" {
		urlPath = "/" + strings.TrimSuffix(out, "index.html")
	}
	return &Page{
		SourcePath: rel, OutputPath: out, URL: urlPath,
		CanonicalURL: "https://tjp.lol" + urlPath, Title: title, Meta: meta,
		ModTime: info.ModTime(), Sections: sections,
	}, warnings, nil
}

func frontmatter(raw []byte) (map[string]any, []byte, error) {
	meta := map[string]any{}
	if !bytes.HasPrefix(raw, []byte("---\n")) && !bytes.HasPrefix(raw, []byte("---\r\n")) {
		return meta, raw, nil
	}
	lines := bytes.SplitAfter(raw, []byte("\n"))
	end, offset := -1, len(lines[0])
	for i := 1; i < len(lines); i++ {
		if strings.TrimSpace(string(lines[i])) == "---" {
			end = i
			break
		}
		offset += len(lines[i])
	}
	if end < 0 {
		return nil, nil, errors.New("unterminated YAML frontmatter")
	}
	front := bytes.Join(lines[1:end], nil)
	if err := yaml.Unmarshal(front, &meta); err != nil {
		return nil, nil, fmt.Errorf("YAML frontmatter: %w", err)
	}
	offset += len(lines[end])
	return meta, raw[offset:], nil
}

func extractTitle(doc ast.Node, source []byte) string {
	for node := doc.FirstChild(); node != nil; node = node.NextSibling() {
		heading, ok := node.(*ast.Heading)
		if !ok || heading.Level != 1 {
			continue
		}
		title := strings.TrimSpace(string(heading.Text(source)))
		doc.RemoveChild(doc, node)
		return title
	}
	return ""
}

func rewriteLinks(root ast.Node) {
	_ = ast.Walk(root, func(node ast.Node, entering bool) (ast.WalkStatus, error) {
		if !entering {
			return ast.WalkContinue, nil
		}
		if link, ok := node.(*ast.Link); ok {
			link.Destination = rewriteDestination(link.Destination)
		}
		return ast.WalkContinue, nil
	})
}

func rewriteDestination(destination []byte) []byte {
	raw := string(destination)
	u, err := url.Parse(raw)
	if err != nil || u.Scheme != "" || u.Host != "" || strings.HasPrefix(raw, "#") || !strings.HasSuffix(strings.ToLower(u.Path), ".md") {
		return destination
	}
	u.Path = u.Path[:len(u.Path)-3] + ".html"
	return []byte(u.String())
}

func renderSections(md goldmark.Markdown, doc ast.Node, source []byte) ([]template.HTML, error) {
	var sections []template.HTML
	for node := doc.FirstChild(); node != nil; {
		next := node.NextSibling()
		doc.RemoveChild(doc, node)
		fragment := ast.NewDocument()
		fragment.AppendChild(fragment, node)
		var out bytes.Buffer
		if err := md.Renderer().Render(&out, source, fragment); err != nil {
			return nil, err
		}
		sections = append(sections, template.HTML(out.String())) // Goldmark raw HTML is disabled.
		node = next
	}
	return sections, nil
}

func (s *site) addPage(page *Page) {
	dir := path.Dir(page.OutputPath)
	if dir == "." {
		dir = ""
	}
	for {
		s.pages[dir] = append(s.pages[dir], page)
		if dir == "" {
			break
		}
		dir = path.Dir(dir)
		if dir == "." {
			dir = ""
		}
	}
}

func (s *site) loadTemplates() error {
	t := template.New("weft").Funcs(templateFuncs())
	for _, src := range s.sources {
		if src.kind != "template" && src.kind != "partial" {
			continue
		}
		contents, err := os.ReadFile(filepath.Join(s.root, filepath.FromSlash(src.path)))
		if err != nil {
			return err
		}
		if _, err := t.New(src.path).Parse(string(contents)); err != nil {
			return fmt.Errorf("parse template %s: %w", src.path, err)
		}
	}
	s.templates = t
	return nil
}

func templateFuncs() template.FuncMap {
	return template.FuncMap{
		"rel":          relativePath,
		"sortPages":    sortPages,
		"filterPages":  filterPages,
		"slicePages":   slicePages,
		"joinSections": joinSections,
		"xml":          func(value any) string { return fmt.Sprint(value) },
		"date":         func(value any, layout string) string { t, _ := asTime(value); return t.Format(layout) },
		"rfc3339":      func(value any) string { t, _ := asTime(value); return t.Format(time.RFC3339) },
	}
}

func relativePath(from, target string) string {
	if strings.HasPrefix(target, "http://") || strings.HasPrefix(target, "https://") || strings.HasPrefix(target, "mailto:") || strings.HasPrefix(target, "#") {
		return target
	}
	fromDir := path.Dir(strings.TrimPrefix(from, "/"))
	if fromDir == "." {
		fromDir = ""
	}
	target = strings.TrimPrefix(target, "/")
	rel, err := filepath.Rel(filepath.FromSlash(fromDir), filepath.FromSlash(target))
	if err != nil {
		return target
	}
	return filepath.ToSlash(rel)
}

func sortPages(pages []*Page, key string) []*Page {
	result := slices.Clone(pages)
	sort.SliceStable(result, func(i, j int) bool {
		if key == "mod_time" {
			return result[i].ModTime.After(result[j].ModTime)
		}
		if key == "title" {
			return strings.ToLower(result[i].Title) < strings.ToLower(result[j].Title)
		}
		a, _ := asTime(result[i].Meta[key])
		b, _ := asTime(result[j].Meta[key])
		return a.After(b)
	})
	return result
}

func filterPages(pages []*Page, key string, value any) []*Page {
	var result []*Page
	for _, page := range pages {
		var actual any
		switch key {
		case "title":
			actual = page.Title
		case "output":
			actual = page.OutputPath
		default:
			actual = page.Meta[key]
		}
		if fmt.Sprint(actual) == fmt.Sprint(value) {
			result = append(result, page)
		}
	}
	return result
}

func slicePages(pages []*Page, start, end int) []*Page {
	if start < 0 {
		start = 0
	}
	if start > len(pages) {
		start = len(pages)
	}
	if end < start {
		end = start
	}
	if end > len(pages) {
		end = len(pages)
	}
	return pages[start:end]
}

func joinSections(sections []template.HTML) template.HTML {
	var out strings.Builder
	for _, section := range sections {
		out.WriteString(string(section))
	}
	return template.HTML(out.String())
}

func asTime(value any) (time.Time, bool) {
	switch v := value.(type) {
	case time.Time:
		return v, true
	case string:
		for _, layout := range []string{time.RFC3339, "2006-01-02"} {
			if parsed, err := time.Parse(layout, v); err == nil {
				return parsed, true
			}
		}
	}
	return time.Time{}, false
}

func (s *site) render() (map[string][]byte, error) {
	outputs := map[string][]byte{}
	for _, src := range s.sources {
		if src.output == "" {
			continue
		}
		var current *Page
		name := src.path
		if src.kind == "markdown" {
			current = src.page
			name = "page"
			if src.output == "weblog" || strings.HasPrefix(src.output, "weblog/") {
				name = "weblog"
			}
			if src.output == "notes" || strings.HasPrefix(src.output, "notes/") {
				name = "note"
			}
		} else {
			current = &Page{SourcePath: src.path, OutputPath: src.output, URL: "/" + src.output, CanonicalURL: "https://tjp.lol/" + src.output, Meta: map[string]any{}}
		}
		if s.templates.Lookup(name) == nil {
			return nil, fmt.Errorf("%s requires missing template %q", src.path, name)
		}
		var out bytes.Buffer
		if err := s.templates.ExecuteTemplate(&out, name, templateData{Page: current, Pages: s.pages}); err != nil {
			return nil, fmt.Errorf("render %s: %w", src.path, err)
		}
		outputs[src.output] = out.Bytes()
	}
	return outputs, nil
}

func (s *site) validate(outputs map[string][]byte) error {
	tracked, err := s.readManifest()
	if err != nil {
		return err
	}
	for output, contents := range outputs {
		if strings.HasSuffix(output, ".xml") {
			var value struct{ XMLName xml.Name }
			if err := xml.Unmarshal(contents, &value); err != nil {
				return fmt.Errorf("invalid XML in %s: %w", output, err)
			}
			if output == "feed.xml" && (value.XMLName.Local != "feed" || value.XMLName.Space != "http://www.w3.org/2005/Atom") {
				return fmt.Errorf("feed.xml is not an Atom 1.0 feed")
			}
		}
		if strings.HasSuffix(output, ".html") {
			if err := s.validateHTMLLinks(output, contents, outputs, tracked); err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *site) validateHTMLLinks(output string, contents []byte, outputs map[string][]byte, tracked map[string]bool) error {
	z := html.NewTokenizer(bytes.NewReader(contents))
	for {
		tokenType := z.Next()
		if tokenType == html.ErrorToken {
			if errors.Is(z.Err(), io.EOF) {
				return nil
			}
			return fmt.Errorf("parse HTML %s: %w", output, z.Err())
		}
		if tokenType != html.StartTagToken && tokenType != html.SelfClosingTagToken {
			continue
		}
		token := z.Token()
		if token.Data != "a" {
			continue
		}
		for _, attr := range token.Attr {
			if attr.Key != "href" {
				continue
			}
			if err := s.validateLink(output, attr.Val, outputs, tracked); err != nil {
				return err
			}
		}
	}
}

func (s *site) validateLink(from, raw string, outputs map[string][]byte, tracked map[string]bool) error {
	u, err := url.Parse(raw)
	if err != nil {
		return fmt.Errorf("invalid link in %s: %q", from, raw)
	}
	if u.Scheme != "" || u.Host != "" || u.Path == "" {
		return nil
	}
	target := u.Path
	if strings.HasPrefix(target, "/") {
		target = strings.TrimPrefix(path.Clean(target), "/")
	} else {
		target = path.Clean(path.Join(path.Dir(from), target))
	}
	if strings.HasSuffix(u.Path, "/") {
		target = path.Join(target, "index.html")
	}
	if target == ".." || strings.HasPrefix(target, "../") {
		return fmt.Errorf("internal link escapes the site root in %s: %s", from, raw)
	}
	if _, ok := outputs[target]; ok {
		return nil
	}
	full := filepath.Join(s.root, filepath.FromSlash(target))
	if info, statErr := os.Stat(full); statErr == nil && !info.IsDir() && !tracked[target] {
		return nil
	}
	indexTarget := path.Join(target, "index.html")
	if info, statErr := os.Stat(filepath.Join(full, "index.html")); statErr == nil && !info.IsDir() && !tracked[indexTarget] {
		return nil
	}
	return fmt.Errorf("broken internal link in %s: %s", from, raw)
}

func (s *site) readManifest() (map[string]bool, error) {
	tracked := map[string]bool{}
	raw, err := os.ReadFile(filepath.Join(s.root, manifestName))
	if errors.Is(err, os.ErrNotExist) {
		return tracked, nil
	}
	if err != nil {
		return nil, err
	}
	var m manifest
	if err := json.Unmarshal(raw, &m); err != nil {
		return nil, fmt.Errorf("read %s: %w", manifestName, err)
	}
	for _, output := range m.Outputs {
		clean := path.Clean(output)
		ext := path.Ext(clean)
		if output == "" || output != clean || strings.Contains(output, `\`) || path.IsAbs(output) || strings.HasPrefix(clean, "../") || ext != ".html" && ext != ".xml" {
			return nil, fmt.Errorf("unsafe path in %s: %q", manifestName, output)
		}
		tracked[output] = true
	}
	return tracked, nil
}

func (s *site) write(outputs map[string][]byte) error {
	tracked, err := s.readManifest()
	if err != nil {
		return err
	}
	for output := range outputs {
		_, statErr := os.Stat(filepath.Join(s.root, filepath.FromSlash(output)))
		if statErr == nil && !tracked[output] {
			return fmt.Errorf("refusing to overwrite untracked output %s", output)
		}
		if statErr != nil && !errors.Is(statErr, os.ErrNotExist) {
			return statErr
		}
	}

	txn, err := os.MkdirTemp(s.root, ".weft-txn-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(txn)
	stage := filepath.Join(txn, "stage")
	backup := filepath.Join(txn, "backup")
	if err := os.MkdirAll(stage, 0o755); err != nil {
		return err
	}
	for output, contents := range outputs {
		name := filepath.Join(stage, filepath.FromSlash(output))
		if err := os.MkdirAll(filepath.Dir(name), 0o755); err != nil {
			return err
		}
		if err := os.WriteFile(name, contents, 0o644); err != nil {
			return err
		}
	}
	paths := make([]string, 0, len(outputs))
	for output := range outputs {
		paths = append(paths, output)
	}
	slices.Sort(paths)
	manifestBytes, _ := json.MarshalIndent(manifest{Outputs: paths}, "", "  ")
	manifestBytes = append(manifestBytes, '\n')
	if err := os.WriteFile(filepath.Join(stage, manifestName), manifestBytes, 0o644); err != nil {
		return err
	}

	var movedOld, installed []string
	rollback := func() {
		for i := len(installed) - 1; i >= 0; i-- {
			_ = os.Remove(filepath.Join(s.root, filepath.FromSlash(installed[i])))
		}
		for i := len(movedOld) - 1; i >= 0; i-- {
			from := filepath.Join(backup, filepath.FromSlash(movedOld[i]))
			to := filepath.Join(s.root, filepath.FromSlash(movedOld[i]))
			_ = os.MkdirAll(filepath.Dir(to), 0o755)
			_ = os.Rename(from, to)
		}
	}
	moveOld := func(rel string) error {
		from := filepath.Join(s.root, filepath.FromSlash(rel))
		if _, err := os.Stat(from); errors.Is(err, os.ErrNotExist) {
			return nil
		} else if err != nil {
			return err
		}
		to := filepath.Join(backup, filepath.FromSlash(rel))
		if err := os.MkdirAll(filepath.Dir(to), 0o755); err != nil {
			return err
		}
		if err := os.Rename(from, to); err != nil {
			return err
		}
		movedOld = append(movedOld, rel)
		return nil
	}
	for old := range tracked {
		if _, wanted := outputs[old]; !wanted {
			if err := moveOld(old); err != nil {
				rollback()
				return err
			}
		}
	}
	for _, rel := range append(paths, manifestName) {
		if err := moveOld(rel); err != nil {
			rollback()
			return err
		}
		to := filepath.Join(s.root, filepath.FromSlash(rel))
		if err := os.MkdirAll(filepath.Dir(to), 0o755); err != nil {
			rollback()
			return err
		}
		if err := os.Rename(filepath.Join(stage, filepath.FromSlash(rel)), to); err != nil {
			rollback()
			return err
		}
		installed = append(installed, rel)
	}
	return nil
}