package main

import (
	"flag"
	"fmt"
	"os"
	"strings"

	"github.com/PuerkitoBio/goquery"
	"golang.org/x/net/html"
)

type Mode string

const (
	ModeQuery   Mode = "query"
	ModeReplace Mode = "replace"
	ModeAppend  Mode = "append"
	ModePrepend Mode = "prepend"
	ModeInvalid Mode = "invalid"

	FilenameStdin string = "-"
)

func main() {
	selector := flag.String("selector", "", "DOM query")
	mode := flag.String("mode", "query", "Operation mode: query, replace, append, prepend")
	filename := flag.String("filename", "", "Path to HTML file")
	writeData := flag.String("wdata", "", "Data to write when in non-query mode")
	flag.Parse()

	//args := flag.Args()

	parsedMode, err := parseMode(*mode)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to parse -mode: %s", err.Error())
		os.Exit(1)
	}

	doc, err := parseFilename(*filename)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to parse -filename: %s", err.Error())
		os.Exit(1)
	}

	switch parsedMode {
	case ModeQuery:
		if *selector == "" {
			// if selector is not specified return whole doc by default
			htmlStr, err := doc.Html()
			if err != nil {
				fmt.Fprintf(os.Stderr, "failed to construct HTML string: %s", err.Error())
				os.Exit(1)
			}
			fmt.Println(htmlStr)
		} else {
			doc.Find(*selector).Each(func(i int, s *goquery.Selection) {
				var buf strings.Builder
				err := html.Render(&buf, s.Get(0))
				if err != nil {
					fmt.Fprintf(os.Stderr, "failed to render node: %s", err.Error())
					os.Exit(1)
				}

				output := buf.String()
				fmt.Println(output)
			})
		}
	case ModeReplace:
		doc.Find(*selector).Each(func(i int, s *goquery.Selection) {
			s.ReplaceWithHtml(*writeData)
		})
	case ModeAppend:
		doc.Find(*selector).Each(func(i int, s *goquery.Selection) {
			s.AppendHtml(*writeData)
		})
	case ModePrepend:
		doc.Find(*selector).Each(func(i int, s *goquery.Selection) {
			s.PrependHtml(*writeData)
		})
	}

	if parsedMode != ModeQuery {
		// in non-query mode, print entire document back with modifications
		htmlStr, err := doc.Html()
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to render node: %s", err.Error())
			os.Exit(1)
		}
		fmt.Println(htmlStr)
	}
}

func parseMode(mode string) (Mode, error) {
	validModes := map[string]any{
		string(ModeQuery):   nil,
		string(ModeReplace): nil,
		string(ModeAppend):  nil,
		string(ModePrepend): nil,
	}

	if _, ok := validModes[mode]; !ok {
		return ModeInvalid, fmt.Errorf("invalid mode: %s", mode)
	}

	return Mode(mode), nil
}

func parseFilename(fileName string) (*goquery.Document, error) {
	var fd *os.File
	var err error

	if fileName == FilenameStdin {
		fd = os.Stdin
	} else {
		fd, err = os.Open(fileName)
		if err != nil {
			return nil, fmt.Errorf("error opening file: %s", err.Error())
		}
		defer fd.Close()
	}

	doc, err := goquery.NewDocumentFromReader(fd)
	if err != nil {
		return nil, fmt.Errorf("error parsing HTML: %s", err.Error())
	}

	return doc, nil
}