package main

import (
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/spf13/cobra"
)

var version = "dev"

var rootCmd = &cobra.Command{
	Use:   "shout",
	Short: "$ shell output tester",
}

func init() {
	rootCmd.AddCommand(testCmd())
	rootCmd.AddCommand(versionCmd())
	rootCmd.AddCommand(exampleCmd())
}

func testCmd() *cobra.Command {
	var (
		update   bool
		keep     bool
		cleanEnv bool
		pathDirs []string
		timeout  string
		verbose  bool
		portFrom int
		parallel bool
	)

	cmd := &cobra.Command{
		Use:   "test [files...]",
		Short: "Run .shout test files",
		Args:  cobra.ArbitraryArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			timeoutDur, err := time.ParseDuration(timeout)
			if err != nil {
				return err
			}

			paths := args
			if len(paths) == 0 {
				paths = []string{"."}
			}

			files, err := findShoutFiles(paths)
			if err != nil {
				return err
			}
			if len(files) == 0 {
				fmt.Fprintln(os.Stderr, "No .shout files found")
				os.Exit(1)
			}

			start := time.Now()
			var results []TestResult
			cwd, _ := os.Getwd()
			nextPort := portFrom

			runOne := func(filePath string, port int) TestResult {
				content, err := os.ReadFile(filePath)
				if err != nil {
					relPath, _ := filepath.Rel(cwd, filePath)
					return TestResult{Path: relPath, Error: err.Error()}
				}

				relPath, _ := filepath.Rel(cwd, filePath)
				parsed, err := parse(relPath, string(content))
				if err != nil {
					return TestResult{Path: relPath, Error: err.Error()}
				}

				// Resolve directives
				envVars := make(map[string]string)
				setupEnvVars := make(map[string]string)
				userEnvVars := make(map[string]string)
				var setupCommands []Command

				for _, d := range parsed.Directives {
					switch d.Type {
					case "setup":
						setupPath := filepath.Join(filepath.Dir(filePath), d.Path)
						setupContent, err := os.ReadFile(setupPath)
						if err != nil {
							return TestResult{Path: parsed.Path, Error: err.Error()}
						}
						setupRelPath, _ := filepath.Rel(cwd, setupPath)
						setupParsed, err := parse(setupRelPath, string(setupContent))
						if err != nil {
							return TestResult{Path: parsed.Path, Error: err.Error()}
						}
						for _, sd := range setupParsed.Directives {
							if sd.Type == "setup" {
								return TestResult{
									Path:  parsed.Path,
									Error: fmt.Sprintf("%s: @setup not allowed in setup files", setupRelPath),
								}
							}
							if sd.Type == "env" {
								setupEnvVars[sd.Key] = sd.Value
							}
						}
						setupCommands = append(setupCommands, setupParsed.Commands...)

					case "env":
						userEnvVars[d.Key] = d.Value
					}
				}

				// Setup env < user env
				for k, v := range setupEnvVars {
					envVars[k] = v
				}
				for k, v := range userEnvVars {
					envVars[k] = v
				}
				if port > 0 {
					if _, ok := userEnvVars["PORT"]; !ok {
						if _, ok := setupEnvVars["PORT"]; !ok {
							envVars["PORT"] = strconv.Itoa(port)
						}
					}
				}

				merged := ShoutFile{
					Path:       parsed.Path,
					Commands:   append(setupCommands, parsed.Commands...),
					Directives: parsed.Directives,
				}

				var onCommand func(Command)
				if verbose {
					onCommand = func(c Command) {
						fmt.Fprintf(os.Stderr, dim("  $ %s\n"), c.Cmd)
					}
				}

				fileResult := runFile(merged, RunOptions{
					CleanEnv:  cleanEnv,
					PathDirs:  pathDirs,
					EnvVars:   envVars,
					Timeout:   timeoutDur,
					OnCommand: onCommand,
				})

				// Check setup commands for failures
				for i := 0; i < len(setupCommands) && i < len(fileResult.Results); i++ {
					r := fileResult.Results[i]
					sc := setupCommands[i]
					if !exitCodeOK(sc.ExitCodeType, sc.ExitCodeValue, r.ExitCode) {
						if keep {
							fmt.Fprintln(os.Stderr, fileResult.TmpDir)
						} else {
							cleanupTmpDir(fileResult.TmpDir)
						}
						return evaluateFile(
							parsed.Path,
							nil,
							fmt.Sprintf("setup command failed (exit %d): $ %s", r.ExitCode, sc.Cmd),
						)
					}
				}

				fileOwnResults := fileResult.Results
				if len(setupCommands) > 0 && len(fileResult.Results) >= len(setupCommands) {
					fileOwnResults = fileResult.Results[len(setupCommands):]
				}

				testResult := evaluateFile(parsed.Path, fileOwnResults, fileResult.Error)

				if update && len(fileOwnResults) > 0 {
					updated := rewriteFile(parsed, fileOwnResults, string(content))
					if updated != string(content) {
						_ = os.WriteFile(filePath, []byte(updated), 0o644)
					}
				}

				if keep {
					fmt.Fprintln(os.Stderr, fileResult.TmpDir)
				} else {
					cleanupTmpDir(fileResult.TmpDir)
				}

				return testResult
			}

			printDots := func(r TestResult) {
				if r.Error != "" {
					fmt.Print(red("F"))
					return
				}
				passed := r.CommandCount - len(r.Failures)
				for i := 0; i < passed; i++ {
					fmt.Print(green("."))
				}
				for i := 0; i < len(r.Failures); i++ {
					fmt.Print(red("F"))
				}
			}

			if parallel {
				allResults := make([]TestResult, len(files))
				var wg sync.WaitGroup
				for idx, f := range files {
					wg.Add(1)
					go func(i int, filePath string, port int) {
						defer wg.Done()
						allResults[i] = runOne(filePath, port)
					}(idx, f, nextPort)
					if nextPort > 0 {
						nextPort++
					}
				}
				wg.Wait()
				for _, r := range allResults {
					printDots(r)
					results = append(results, r)
				}
				fmt.Println()
			} else {
				for _, filePath := range files {
					port := 0
					if nextPort > 0 {
						port = nextPort
						nextPort++
					}
					r := runOne(filePath, port)
					printDots(r)
					results = append(results, r)
				}
				fmt.Println()
			}

			// Print failures
			var failures []TestResult
			for _, r := range results {
				if !r.Passed {
					failures = append(failures, r)
				}
			}
			if len(failures) > 0 {
				fmt.Println()
				for _, f := range failures {
					fmt.Println(formatFailure(f))
					fmt.Println()
				}
			}

			elapsed := time.Since(start)
			fmt.Println(formatSummary(results, elapsed))

			if len(failures) > 0 {
				os.Exit(1)
			}
			return nil
		},
	}

	cmd.Flags().BoolVarP(&update, "update", "u", false, "Rewrite expected output in-place with actual output")
	cmd.Flags().BoolVarP(&keep, "keep", "k", false, "Keep temp directories after run")
	cmd.Flags().BoolVar(&cleanEnv, "clean-env", false, "Start with empty environment")
	cmd.Flags().StringArrayVar(&pathDirs, "path", nil, "Prepend <path> to PATH (repeatable)")
	cmd.Flags().StringVar(&timeout, "timeout", "10s", "Per-command timeout")
	cmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "Print each command as it runs")
	cmd.Flags().IntVar(&portFrom, "port-from", 0, "Auto-assign $PORT starting from <n>")
	cmd.Flags().BoolVar(&parallel, "parallel", false, "Run files in parallel")

	return cmd
}

func versionCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "version",
		Short: "Print the version",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Println(version)
		},
	}
}

func exampleCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "example",
		Short: "Print an example .shout file",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Print(`# Example .shout file
$ echo hello
hello

$ echo "one"; echo "two"; echo "three"
one
...
three

$ cat nonexistent
cat: nonexistent: ...
[1]

$ true
[0]
`)
		},
	}
}

func findShoutFiles(paths []string) ([]string, error) {
	var files []string

	for _, p := range paths {
		abs, err := filepath.Abs(p)
		if err != nil {
			continue
		}

		info, err := os.Stat(abs)
		if err != nil {
			if strings.HasSuffix(abs, ".shout") {
				files = append(files, abs)
			}
			continue
		}

		if info.IsDir() {
			_ = filepath.WalkDir(abs, func(path string, d fs.DirEntry, err error) error {
				if err != nil {
					return nil
				}
				if !d.IsDir() && strings.HasSuffix(path, ".shout") {
					files = append(files, path)
				}
				return nil
			})
		} else if strings.HasSuffix(abs, ".shout") {
			files = append(files, abs)
		}
	}

	sort.Strings(files)
	return files, nil
}