package main

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"

	"github.com/marzeq/qk/shared"
)

var compilerVersion = "(in development)"

type OutputType int

const (
	OutputUnspecified OutputType = iota
	OutputExecutable
	OutputObject
	OutputSharedLib
	OutputWebAssembly
	OutputStaticLibrary
)

type OptimisationLevel string

const (
	OptLevel0       OptimisationLevel = "0"
	OptLevel1       OptimisationLevel = "1"
	OptLevel2       OptimisationLevel = "2"
	OptLevel3       OptimisationLevel = "3"
	OptLevelSize    OptimisationLevel = "s"
	OptLevelSizeMax OptimisationLevel = "z"
	OptLevelFast    OptimisationLevel = "fast"
	OptLevelDebug   OptimisationLevel = "g"
)

type WarningMode string

const (
	WarningModeShow  WarningMode = "show"
	WarningModeOff   WarningMode = "off"
	WarningModeError WarningMode = "error"
)

func parseWarningMode(value string) (WarningMode, error) {
	switch WarningMode(value) {
	case WarningModeShow, WarningModeOff, WarningModeError:
		return WarningMode(value), nil
	default:
		return "", fmt.Errorf("invalid warning mode %q: expected show, off, or error", value)
	}
}

type Args struct {
	baseDir      string
	packageRoot  string
	packagePaths []string
	sourceMounts []sourceMount
	manifestPath string
	manifestData string
	file         string
	packageArg   string
	programArgs  []string
	installPath  string
	output       string
	mainModule   string
	outputName   string
	optLevel     OptimisationLevel
	verbose      bool
	quietLink    bool
	debug        bool
	warningMode  WarningMode
	warningModes map[string]WarningMode
	dumpIR       bool
	dumpLLVM     bool
	dumpAsm      bool
	keepBuildDir bool
	static       bool
	release      bool
	libraryPath  string
	noStdlib     bool
	noEmit       bool
	target       string
	sysroot      string
	cpu          string
	features     string
	targetABI    string
	relocation   string
	codeModel    string
	outputType   OutputType
	linkArgs     []string
	libs         []string
	libraryPaths []string
	run          bool
	install      bool
	cpuProfile   string
}

func parseOptLevel(level string) (OptimisationLevel, error) {
	switch level {
	case "g":
		return OptLevelDebug, nil
	case "s":
		return OptLevelSize, nil
	case "z":
		return OptLevelSizeMax, nil
	case "fast":
		return OptLevelFast, nil
	default:
		n, err := strconv.Atoi(level)
		if err != nil || n < 0 {
			return "", fmt.Errorf("invalid optimisation level: %s", level)
		}
		if n > 3 {
			fmt.Fprintf(os.Stderr, "warning: optimisation level %s is equivalent to -O3\n", level)
			n = 3
		}
		return OptimisationLevel(strconv.Itoa(n)), nil
	}
}

func parseOutputType(value string) (OutputType, error) {
	switch value {
	case "exe", "executable", ".exe":
		return OutputExecutable, nil
	case "obj", "object", ".o", ".obj":
		return OutputObject, nil
	case "lib", "staticlib", "archive", ".a", ".lib":
		return OutputStaticLibrary, nil
	case "so", "shared", "sharedlib", ".so", ".dll", ".dylib":
		return OutputSharedLib, nil
	case "wasm", ".wasm":
		return OutputWebAssembly, nil
	default:
		return OutputUnspecified, fmt.Errorf("unknown output type: %s", value)
	}
}

type argumentParser struct {
	args  *Args
	input []string
	index int
}

func newArgumentParser(input []string) *argumentParser {
	return &argumentParser{
		args: &Args{
			optLevel:     OptLevel2,
			outputType:   OutputUnspecified,
			warningMode:  WarningModeShow,
			warningModes: make(map[string]WarningMode),
		},
		input: input,
	}
}

func (p *argumentParser) current() string {
	return p.input[p.index]
}

func (p *argumentParser) nextValue(option string) (string, error) {
	p.index++
	if p.index >= len(p.input) {
		return "", fmt.Errorf("expected value after %s", option)
	}
	value := p.input[p.index]
	p.index++
	return value, nil
}

func (p *argumentParser) gluedOrNextValue(prefix string) (string, error) {
	option := p.current()
	if len(option) > len(prefix) {
		p.index++
		return option[len(prefix):], nil
	}
	return p.nextValue(prefix)
}

func (p *argumentParser) parseSplitArgs(option string, target *[]string) error {
	value, err := p.nextValue(option)
	if err != nil {
		return err
	}
	*target = append(*target, strings.Fields(value)...)
	return nil
}

func (p *argumentParser) parse() (*Args, error) {
	if len(p.input) == 0 {
		return nil, fmt.Errorf("expected build, run, or install command")
	}
	switch p.input[0] {
	case "build":
	case "run":
		p.args.run = true
	case "install":
		p.args.install = true
		p.args.release = true
	case "-h", "--help":
		printUsage()
		os.Exit(0)
	case "-v", "--version":
		printVersion()
		os.Exit(0)
	default:
		return nil, fmt.Errorf("unknown command %q: expected build, run, or install", p.input[0])
	}
	p.index = 1
	for p.index < len(p.input) {
		if p.args.run && p.args.packageArg != "" {
			p.args.programArgs = append(p.args.programArgs, p.input[p.index:]...)
			p.index = len(p.input)
			break
		}
		if err := p.parseCurrent(); err != nil {
			return nil, err
		}
	}
	if p.args.install {
		if p.args.packageArg == "" {
			return nil, fmt.Errorf("install requires a package path and an install directory")
		}
		if p.args.installPath == "" {
			return nil, fmt.Errorf("install requires an install directory after the package path")
		}
	}

	if err := finaliseArgs(p.args); err != nil {
		return nil, err
	}
	return p.args, nil
}

func (p *argumentParser) parseCurrent() error {
	tok := p.current()

	switch {
	case tok == "-o":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		p.args.output = value

	case tok == "-t":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		outputType, err := parseOutputType(value)
		if err != nil {
			return err
		}
		p.args.outputType = outputType

	case strings.HasPrefix(tok, "-O"):
		value, err := p.gluedOrNextValue("-O")
		if err != nil {
			return err
		}
		level, err := parseOptLevel(value)
		if err != nil {
			return err
		}
		p.args.optLevel = level

	case tok == "-v":
		p.args.verbose = true
		p.index++

	case tok == "-d":
		p.args.debug = true
		p.index++

	case tok == "-warn":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		mode, err := parseWarningMode(value)
		if err != nil {
			return err
		}
		p.args.warningMode = mode

	case strings.HasPrefix(tok, "-warn-"):
		warningType := strings.TrimPrefix(tok, "-warn-")
		switch warningType {
		case string(shared.WarningUnusedVariable), string(shared.WarningUnusedParameter):
		default:
			return fmt.Errorf("unknown warning type %q", warningType)
		}
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		mode, err := parseWarningMode(value)
		if err != nil {
			return err
		}
		p.args.warningModes[warningType] = mode

	case tok == "-no-emit":
		p.args.noEmit = true
		p.index++

	case tok == "-cpuprofile":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		p.args.cpuProfile = value

	case tok == "-dump-ir":
		p.args.dumpIR = true
		p.index++

	case tok == "-dump-llvm":
		p.args.dumpLLVM = true
		p.index++

	case tok == "-dump-asm":
		p.args.dumpAsm = true
		p.index++

	case tok == "-keep-build-dir":
		p.args.keepBuildDir = true
		p.index++

	case tok == "-static":
		p.args.static = true
		p.index++

	case tok == "-release":
		p.args.release = true
		p.index++

	case tok == "-nostdlib":
		p.args.noStdlib = true
		p.index++

	case tok == "-stdlib":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		p.args.libraryPath = value

	case tok == "-target":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		p.args.target = value

	case tok == "-sysroot":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		p.args.sysroot = value

	case tok == "-cpu":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		p.args.cpu = value

	case tok == "-features":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		p.args.features = value

	case tok == "-target-abi":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		p.args.targetABI = value

	case tok == "-relocation-model":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		switch value {
		case "default", "static", "pic", "dynamic-no-pic":
			p.args.relocation = value
		default:
			return fmt.Errorf("invalid relocation model: %s", value)
		}

	case tok == "-code-model":
		value, err := p.nextValue(tok)
		if err != nil {
			return err
		}
		switch value {
		case "default", "tiny", "small", "kernel", "medium", "large":
			p.args.codeModel = value
		default:
			return fmt.Errorf("invalid code model: %s", value)
		}

	case tok == "-Xlink":
		if err := p.parseSplitArgs(tok, &p.args.linkArgs); err != nil {
			return err
		}

	case strings.HasPrefix(tok, "-I"):
		value, err := p.gluedOrNextValue("-I")
		if err != nil {
			return err
		}
		p.args.packagePaths = append(p.args.packagePaths, value)

	case strings.HasPrefix(tok, "-l"):
		value, err := p.gluedOrNextValue("-l")
		if err != nil {
			return err
		}
		p.args.libs = append(p.args.libs, value)

	case strings.HasPrefix(tok, "-L"):
		value, err := p.gluedOrNextValue("-L")
		if err != nil {
			return err
		}
		p.args.libraryPaths = append(p.args.libraryPaths, value)

	case tok == "-h" || tok == "--help":
		printUsage()
		os.Exit(0)

	default:
		if strings.HasPrefix(tok, "-") {
			return fmt.Errorf("unknown argument: %s", tok)
		}
		if p.args.install && p.args.packageArg != "" && p.args.installPath == "" {
			p.args.installPath = tok
			p.index++
			return nil
		}
		if p.args.packageArg != "" {
			return fmt.Errorf("multiple package arguments specified")
		}
		p.args.packageArg = tok
		p.index++
	}

	return nil
}

func parseArgs() (*Args, error) {
	return newArgumentParser(os.Args[1:]).parse()
}

func printUsage() {
	name := filepath.Base(os.Args[0])
	fmt.Printf(`QK compiler

Usage:
  %s build [options] [package]
  %s run [options] [package] [arguments...]
  %s install [options] <package> <directory>
  %s --help
  %s --version

Commands:
  build      Compile a package (the current directory by default)
  run        Compile and run an executable package
  install    Compile an executable in release mode and place it in directory

A package may be a directory or one explicit .qk source file. Options for run
must appear before the package; all following values are passed to the program.

Build options:
  -o <file>             Set the output file name
  -t <type>             Set output type: exe, obj, lib, so, or wasm
  -O <level>            Set optimisation: 0, 1, 2, 3, s, z, fast, or g
  -release              Select .Release for the compile-time ReleaseMode value
  -static               Prefer static libraries when linking
  -no-emit              Check and compile without writing an output file

Diagnostics:
  -warn <mode>          Set warnings to show, off, or error (default: show)
  -warn-<name> <mode>   Override one warning category
  -d                    Enable compiler debug output
  -v                    Print verbose compiler and linker output
  -dump-ir              Print target-independent QK IR
  -dump-llvm            Print generated LLVM IR
  -dump-asm             Print generated assembly
  -keep-build-dir       Preserve temporary build directories
  -cpuprofile <file>    Write a compiler CPU profile

Packages and libraries:
  -nostdlib             Do not load the installed QK library tree
  -stdlib <dir>         Use an external QK library root
  -I <dir>              Add a package search root (repeatable)
  -L <dir>              Add a library search path (repeatable)
  -l <lib>              Link a library (repeatable)
  -Xlink <args>         Pass space-separated arguments to the linker

Target options:
  -target <triple>      Set the LLVM target triple
  -sysroot <path>       Set the target sysroot
  -cpu <name>           Set the LLVM target CPU
  -features <list>      Set LLVM target features
  -target-abi <name>    Set the target-specific ABI
  -relocation-model <model>
                        Set default, static, pic, or dynamic-no-pic
  -code-model <model>   Set default, tiny, small, kernel, medium, or large
`, name, name, name, name, name)
}

func printVersion() {
	fmt.Printf("qk compiler version %s\n", compilerVersion)
}

func finaliseArgs(args *Args) error {
	if args.target == "" {
		args.target = defaultTargetForHost(runtime.GOOS)
	}
	if args.packageArg == "" {
		args.packageArg = "."
	}
	info, err := os.Stat(args.packageArg)
	if err != nil {
		return fmt.Errorf("package path does not exist: %s", args.packageArg)
	}
	abs, err := filepath.Abs(args.packageArg)
	if err != nil {
		return fmt.Errorf("failed to get absolute package path: %v", err)
	}
	if info.IsDir() {
		args.baseDir = abs
		args.outputName = filepath.Base(abs)
	} else {
		if !info.Mode().IsRegular() || !strings.EqualFold(filepath.Ext(abs), ".qk") {
			return fmt.Errorf("package argument must be a directory or .qk source file: %s", args.packageArg)
		}
		args.file = abs
		args.baseDir = filepath.Dir(abs)
		args.outputName = strings.TrimSuffix(filepath.Base(abs), filepath.Ext(abs))
	}
	workingDir, err := os.Getwd()
	if err != nil {
		return fmt.Errorf("failed to get working directory: %v", err)
	}
	manifest, err := findProjectManifest(args.baseDir)
	if err != nil {
		return err
	}
	args.packageRoot = workingDir
	if manifest != nil {
		args.packageRoot = manifest.Root
		args.sourceMounts = manifest.Mounts
		args.manifestPath = manifest.Path
		args.manifestData = manifest.Source
	} else if !pathWithin(args.baseDir, workingDir) {
		args.packageRoot = args.baseDir
	}
	args.mainModule = packagePathFromDirectory(args.packageRoot, args.baseDir)
	if args.install {
		if args.output != "" {
			return fmt.Errorf("-o cannot be used with install")
		}
		if args.outputType != OutputUnspecified && args.outputType != OutputExecutable {
			return fmt.Errorf("install only supports executable output")
		}
		info, err := os.Stat(args.installPath)
		if err != nil || !info.IsDir() {
			return fmt.Errorf("install path is not a directory: %s", args.installPath)
		}
		installPath, err := filepath.Abs(args.installPath)
		if err != nil {
			return fmt.Errorf("failed to get absolute install path: %v", err)
		}
		args.installPath = installPath
		args.outputType = OutputExecutable
		args.output = filepath.Join(installPath, defaultExecutableName(args.outputName, args.target))
	}
	for i, path := range args.packagePaths {
		info, err := os.Stat(path)
		if err != nil || !info.IsDir() {
			return fmt.Errorf("package search root is not a directory: %s", path)
		}
		abs, err := filepath.Abs(path)
		if err != nil {
			return fmt.Errorf("failed to get absolute package search root: %v", err)
		}
		args.packagePaths[i] = abs
	}

	if args.sysroot != "" {
		if _, err := os.Stat(args.sysroot); os.IsNotExist(err) {
			return fmt.Errorf("sysroot path does not exist: %s", args.sysroot)
		}
		abs, err := filepath.Abs(args.sysroot)
		if err != nil {
			return fmt.Errorf("failed to get absolute path of sysroot: %v", err)
		}
		args.sysroot = abs
	}
	if args.libraryPath != "" {
		if args.noStdlib {
			return fmt.Errorf("-nostdlib and -stdlib cannot be used together")
		}
		info, err := os.Stat(args.libraryPath)
		if err != nil || !info.IsDir() {
			return fmt.Errorf("QK library path is not a directory: %s", args.libraryPath)
		}
		abs, err := filepath.Abs(args.libraryPath)
		if err != nil {
			return fmt.Errorf("failed to get absolute QK library path: %v", err)
		}
		args.libraryPath = abs
	}

	return nil
}

func finaliseOutputArgs(args *Args) error {
	if args.output == "" {
		switch args.outputType {
		case OutputUnspecified:
			args.output = defaultExecutableName(args.outputName, args.target)
			args.outputType = OutputExecutable
		case OutputExecutable:
			args.output = defaultExecutableName(args.outputName, args.target)
		case OutputObject:
			args.output = defaultObjectName(args.outputName, args.target)
		case OutputSharedLib:
			args.output = defaultSharedLibraryName(args.outputName, args.target)
		case OutputWebAssembly:
			args.output = args.outputName + ".wasm"
		case OutputStaticLibrary:
			args.output = defaultStaticLibraryName(args.outputName, args.target)
		}
	} else {
		switch args.outputType {
		case OutputUnspecified:
			ext := strings.ToLower(filepath.Ext(args.output))
			switch ext {
			case ".o", ".obj":
				args.outputType = OutputObject
			case ".a", ".lib":
				args.outputType = OutputStaticLibrary
			case ".wasm":
				args.outputType = OutputWebAssembly
			case ".so", ".dll", ".dylib":
				args.outputType = OutputSharedLib
			case "", ".exe":
				args.outputType = OutputExecutable
			default:
				return fmt.Errorf("cannot infer output type from extension: %s", ext)
			}
		}
	}
	if args.outputType == OutputWebAssembly && !targetIsWebAssembly(args.target) {
		return fmt.Errorf("cannot use WebAssembly output for non-WebAssembly target %q", effectiveTargetName(args.target))
	}

	if args.run && args.outputType != OutputExecutable {
		return fmt.Errorf("cannot run non-executable output")
	}

	return nil
}

func defaultExecutableName(module, target string) string {
	if targetIsWindows(target) {
		return module + ".exe"
	}
	return module
}

func defaultObjectName(module, target string) string {
	if targetIsWindows(target) {
		return module + ".obj"
	}
	return module + ".o"
}

func defaultStaticLibraryName(module, target string) string {
	if targetIsWindowsMSVC(target) {
		return module + ".lib"
	}
	return "lib" + module + ".a"
}

func defaultSharedLibraryName(module, target string) string {
	switch {
	case targetIsWindows(target):
		return module + ".dll"
	case targetIsApple(target):
		return "lib" + module + ".dylib"
	default:
		return "lib" + module + ".so"
	}
}

func effectiveTargetName(target string) string {
	if target != "" {
		return strings.ToLower(target)
	}
	if target = defaultTargetForHost(runtime.GOOS); target != "" {
		return target
	}
	return runtime.GOARCH + "-" + runtime.GOOS
}

func defaultTargetForHost(goos string) string {
	if goos == "windows" {
		return "x86_64-pc-windows-msvc"
	}
	return ""
}

func targetIsWindows(target string) bool {
	target = effectiveTargetName(target)
	return strings.Contains(target, "windows") || strings.Contains(target, "mingw") || strings.Contains(target, "msvc")
}

func targetIsApple(target string) bool {
	target = effectiveTargetName(target)
	return strings.Contains(target, "darwin") || strings.Contains(target, "apple") || strings.Contains(target, "macos") || strings.Contains(target, "ios")
}

func targetIsWebAssembly(target string) bool {
	target = effectiveTargetName(target)
	return strings.HasPrefix(target, "wasm32-") || strings.HasPrefix(target, "wasm64-")
}