package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"c65gm/internal/commands"
	"c65gm/internal/compiler"
	"c65gm/internal/preproc"
)

// c65gm - A 6502 Cross-Compiler for the ACME Cross-Assembler
// Copyright (C) 1999, 2025 Mattias Hansson
// Distributed under GPL.

// ANSI color codes for error messages
const (
	colorRed    = "\033[31m"
	colorYellow = "\033[33m"
	colorReset  = "\033[0m"
)

func main() {
	fmt.Println("c65gm - A 6502 Cross-Compiler for the ACME Cross-Assembler.")
	fmt.Println("Copyright (C) 1999, 2025 Mattias Hansson. v1.0.0")
	fmt.Println("Distributed under GPL.")
	fmt.Println()

	// Check if we have any arguments
	if len(os.Args) < 2 {
		printUsage()
		os.Exit(1)
	}

	// Check for help flags
	if os.Args[1] == "-h" || os.Args[1] == "--help" {
		printUsage()
		return
	}
	
	// Check if first arg is a subcommand
	subcommand := os.Args[1]
	if subcommand == "build" || subcommand == "compile" || subcommand == "help" {
		switch subcommand {
		case "build":
			runBuildCommand(os.Args[2:])
		case "compile":
			runCompileCommand(os.Args[2:])
		case "help":
			printUsage()
		}
		return
	}
	
	// Default mode: treat as build command with implicit arguments
	// Parse arguments flexibly
	var inputFile, outputFile string
	args := os.Args[1:]
	
	for i := 0; i < len(args); i++ {
		arg := args[i]
		
		if arg == "-i" || arg == "-in" {
			if i+1 < len(args) {
				inputFile = args[i+1]
				i++ // Skip next arg
			} else {
				fmt.Fprintln(os.Stderr, "Error: -i flag requires a filename")
				printUsage()
				os.Exit(1)
			}
		} else if arg == "-o" || arg == "-out" {
			if i+1 < len(args) {
				outputFile = args[i+1]
				i++ // Skip next arg
			} else {
				fmt.Fprintln(os.Stderr, "Error: -o flag requires a filename")
				printUsage()
				os.Exit(1)
			}
		} else if !strings.HasPrefix(arg, "-") && inputFile == "" {
			// First non-flag argument is the input file
			inputFile = arg
		}
	}
	
	if inputFile == "" {
		fmt.Fprintln(os.Stderr, "Error: No input file specified")
		printUsage()
		os.Exit(1)
	}
	
	// Default output based on input
	if outputFile == "" {
		base := strings.TrimSuffix(filepath.Base(inputFile), filepath.Ext(inputFile))
		outputFile = base + ".prg" // Default to build mode (.prg)
	}
	
	// Determine mode by output extension
	if strings.HasSuffix(strings.ToLower(outputFile), ".prg") {
		// Build mode (compile + assemble)
		if err := build(inputFile, outputFile, false, false); err != nil {
			handleError(err)
		}
		fmt.Println("Build successful.")
	} else {
		// Compile mode (assembly only)
		if err := compileOnly(inputFile, outputFile); err != nil {
			handleError(err)
		}
		fmt.Println("Compilation successful.")
	}
}

func compileOnly(inFile, outFile string) error {
	// Preprocess
	lines, pragma, err := preproc.PreProcess(inFile)
	if err != nil {
		return fmt.Errorf("preprocessing failed: %w", err)
	}

	fmt.Printf("Preprocessed %d lines\n", len(lines))

	// Create compiler and register commands
	comp := compiler.NewCompiler(pragma)
	registerCommands(comp)

	// Compile
	asmLines, err := comp.Compile(lines)
	if err != nil {
		return fmt.Errorf("compilation failed: %w", err)
	}

	fmt.Printf("Generated %d lines of assembly\n", len(asmLines))

	// Write output
	if err := writeOutput(outFile, asmLines); err != nil {
		return fmt.Errorf("failed to write output: %w", err)
	}

	return nil
}

func registerCommands(comp *compiler.Compiler) {
	// Register all command handlers here
	// This is the single place where all commands are wired up

	comp.Registry().Register(&commands.ByteCommand{})
	comp.Registry().Register(&commands.WordCommand{})
	comp.Registry().Register(&commands.AddCommand{})
	comp.Registry().Register(&commands.AndCommand{})
	comp.Registry().Register(&commands.OrCommand{})
	comp.Registry().Register(&commands.XorCommand{})
	comp.Registry().Register(&commands.SubtractCommand{})
	comp.Registry().Register(&commands.LetCommand{})
	comp.Registry().Register(&commands.IfCommand{})
	comp.Registry().Register(&commands.ElseCommand{})
	comp.Registry().Register(&commands.EndIfCommand{})
	comp.Registry().Register(&commands.WhileCommand{})
	comp.Registry().Register(&commands.BreakCommand{})
	comp.Registry().Register(&commands.WendCommand{})
	comp.Registry().Register(&commands.FuncCommand{})
	comp.Registry().Register(commands.NewCallCommand(comp.Context().FunctionHandler))
	comp.Registry().Register(&commands.FendCommand{})
	comp.Registry().Register(&commands.IncrCommand{})
	comp.Registry().Register(&commands.DecrCommand{})
	comp.Registry().Register(&commands.GotoCommand{})
	comp.Registry().Register(&commands.LabelCommand{})
	comp.Registry().Register(&commands.OriginCommand{})
	comp.Registry().Register(&commands.PointerCommand{})
	comp.Registry().Register(&commands.PeekCommand{})
	comp.Registry().Register(&commands.PeekWCommand{})
	comp.Registry().Register(&commands.PokeCommand{})
	comp.Registry().Register(&commands.PokeWCommand{})
	comp.Registry().Register(&commands.SubEndCommand{})
	comp.Registry().Register(&commands.GosubCommand{})
	comp.Registry().Register(&commands.ForCommand{})
	comp.Registry().Register(&commands.NextCommand{})
	comp.Registry().Register(&commands.SwitchCommand{})
	comp.Registry().Register(&commands.CaseCommand{})
	comp.Registry().Register(&commands.DefaultCommand{})
	comp.Registry().Register(&commands.EndSwitchCommand{})
	comp.Registry().Register(&commands.MacroCommand{})
	comp.Registry().Register(&commands.ShiftLCommand{})
	comp.Registry().Register(&commands.ShiftRCommand{})
}

func writeOutput(filename string, lines []string) error {
	return os.WriteFile(filename, []byte(strings.Join(lines, "\n")+"\n"), 0644)
}

func handleError(err error) {
	if _, ok := err.(preproc.HaltError); ok {
		fmt.Println("Halted by #HALT directive")
		os.Exit(2)
	}
	_, _ = fmt.Fprintf(os.Stderr, "Error: %v\n", err)
	os.Exit(1)
}

func printUsage() {
	fmt.Println("Usage:")
	fmt.Println("  c65gm build -i <input.c65> [-o <output.prg>] [--keep-asm] [--no-cbm]")
	fmt.Println("  c65gm compile -i <input.c65> [-o <output.asm>]")
	fmt.Println("  c65gm <input.c65> [-o <output>] (smart default mode)")
	fmt.Println("  c65gm -in <input.c65> -out <output.asm> (legacy mode)")
	fmt.Println()
	fmt.Println("Commands:")
	fmt.Println("  build     Compile and assemble to .prg file (default)")
	fmt.Println("  compile   Compile to .asm file only")
	fmt.Println("  help      Show this help")
	fmt.Println()
	fmt.Println("Options:")
	fmt.Println("  -i, --input   Input .c65 file (required)")
	fmt.Println("  -o, --output  Output file (default: <input>.prg for build, <input>.asm for compile)")
	fmt.Println("  --keep-asm    Keep intermediate assembly file (build command only)")
	fmt.Println("  --no-cbm      Don't add -f cbm flag to ACME (build command only)")
}

func runBuildCommand(args []string) {
	buildCmd := flag.NewFlagSet("build", flag.ExitOnError)
	input := buildCmd.String("i", "", "input .c65 file (required)")
	output := buildCmd.String("o", "", "output .prg file (default: <input>.prg)")
	keepAsm := buildCmd.Bool("keep-asm", false, "keep intermediate assembly file")
	noCBM := buildCmd.Bool("no-cbm", false, "don't add -f cbm flag to ACME")
	
	if err := buildCmd.Parse(args); err != nil {
		fmt.Fprintf(os.Stderr, "Error parsing flags: %v\n", err)
		os.Exit(1)
	}
	
	if *input == "" {
		fmt.Fprintln(os.Stderr, "Error: -i flag is required")
		buildCmd.Usage()
		os.Exit(1)
	}
	
	// Default output filename
	if *output == "" {
		base := strings.TrimSuffix(filepath.Base(*input), filepath.Ext(*input))
		*output = base + ".prg"
	}
	
	// Ensure output has .prg extension
	if !strings.HasSuffix(strings.ToLower(*output), ".prg") {
		*output = *output + ".prg"
	}
	
	if err := build(*input, *output, *keepAsm, *noCBM); err != nil {
		handleError(err)
	}
	fmt.Println("Build successful.")
}

func runCompileCommand(args []string) {
	compileCmd := flag.NewFlagSet("compile", flag.ExitOnError)
	input := compileCmd.String("i", "", "input .c65 file (required)")
	output := compileCmd.String("o", "", "output .asm file (default: <input>.asm)")
	
	if err := compileCmd.Parse(args); err != nil {
		fmt.Fprintf(os.Stderr, "Error parsing flags: %v\n", err)
		os.Exit(1)
	}
	
	if *input == "" {
		fmt.Fprintln(os.Stderr, "Error: -i flag is required")
		compileCmd.Usage()
		os.Exit(1)
	}
	
	// Default output filename
	if *output == "" {
		base := strings.TrimSuffix(filepath.Base(*input), filepath.Ext(*input))
		*output = base + ".asm"
	}
	
	// Ensure output has .asm extension
	if !strings.HasSuffix(strings.ToLower(*output), ".asm") {
		*output = *output + ".asm"
	}
	
	if err := compileOnly(*input, *output); err != nil {
		handleError(err)
	}
	fmt.Println("Compilation successful.")
}

func build(inputFile, outputFile string, keepAsm, noCBM bool) error {
	// Check if ACME is available before starting compilation
	if err := checkACMEAvailable(); err != nil {
		return err
	}
	
	// Assembly file sits in current directory for easy inspection
	base := strings.TrimSuffix(filepath.Base(inputFile), filepath.Ext(inputFile))
	asmFile := base + ".asm"
	
	// Compile to assembly
	if err := compileOnly(inputFile, asmFile); err != nil {
		return fmt.Errorf("compilation failed: %w", err)
	}
	
	// Run ACME assembler
	if err := runACME(asmFile, outputFile, noCBM); err != nil {
		return fmt.Errorf("assembly failed: %w", err)
	}
	
	// Remove assembly file unless requested to keep it
	if !keepAsm {
		if err := os.Remove(asmFile); err != nil {
			return fmt.Errorf("failed to remove assembly file: %w", err)
		}
	} else {
		fmt.Printf("Intermediate assembly file kept: %s\n", asmFile)
	}
	
	return nil
}

func runACME(asmFile, outputFile string, noCBM bool) error {
	// Check if ACME is available
	acmePath, err := exec.LookPath("acme")
	if err != nil {
		return fmt.Errorf("acme assembler not found in PATH")
	}
	
	// Build ACME command arguments
	args := []string{"-o", outputFile}
	if !noCBM {
		args = append(args, "-f", "cbm")
	}
	args = append(args, asmFile)
	
	fmt.Printf("Running ACME assembler: %s %s\n", acmePath, strings.Join(args, " "))
	
	cmd := exec.Command(acmePath, args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("acme failed: %w", err)
	}
	
	fmt.Printf("Assembled to: %s\n", outputFile)
	return nil
}

// detectPackageManager detects available package managers and returns platform-specific installation instructions
func detectPackageManager() (string, string) {
	// Platform detection
	os := runtime.GOOS
	
	// Helper function to check if a command exists
	commandExists := func(cmd string) bool {
		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
		defer cancel()
		
		// Try "command -v" first (POSIX standard)
		checkCmd := exec.CommandContext(ctx, "command", "-v", cmd)
		if err := checkCmd.Run(); err == nil {
			return true
		}
		
		// Fallback to "which" (common but not POSIX)
		checkCmd = exec.CommandContext(ctx, "which", cmd)
		if err := checkCmd.Run(); err == nil {
			return true
		}
		
		return false
	}
	
	switch os {
	case "linux":
		// Linux package manager detection (priority order)
		if commandExists("apt") {
			return "Ubuntu/Debian (apt)", "sudo apt install acme"
		}
		if commandExists("dnf") {
			return "Fedora/RHEL (dnf)", "sudo dnf install acme"
		}
		if commandExists("yum") {
			return "RHEL/CentOS (yum)", "sudo yum install acme"
		}
		if commandExists("pacman") {
			return "Arch (pacman)", "sudo pacman -S acme"
		}
		if commandExists("zypper") {
			return "openSUSE (zypper)", "sudo zypper install acme"
		}
		// Generic Linux fallback
		return "Linux", "Download from releases or use your package manager"
		
	case "darwin": // macOS
		if commandExists("brew") {
			return "macOS (Homebrew)", "brew install acme"
		}
		return "macOS", "brew install acme  # Install Homebrew first, or download binary from releases"
		
	case "windows":
		// Windows package manager detection
		if commandExists("choco") {
			return "Windows (Chocolatey)", "choco install acme"
		}
		if commandExists("scoop") {
			return "Windows (Scoop)", "scoop install acme"
		}
		// Check for winget (Windows Package Manager)
		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
		defer cancel()
		checkCmd := exec.CommandContext(ctx, "cmd", "/c", "where", "winget")
		if err := checkCmd.Run(); err == nil {
			return "Windows (winget)", "winget install acme"
		}
		return "Windows", "Download binary from: https://github.com/meonwax/acme/releases"
		
	default:
		return "Unknown OS", "Download from: https://github.com/meonwax/acme"
	}
}

// checkACMEAvailable checks if ACME is installed and provides helpful instructions if not
func checkACMEAvailable() error {
	_, err := exec.LookPath("acme")
	if err != nil {
		platform, installCmd := detectPackageManager()
		
		var message strings.Builder
		message.WriteString(fmt.Sprintf("%sACME assembler not found in PATH%s\n\n", colorRed, colorReset))
		message.WriteString("c65gm requires ACME (Cross-Assembler) to create executable files.\n\n")
		
		message.WriteString(fmt.Sprintf("%sInstallation options for %s:%s\n", colorYellow, platform, colorReset))
		message.WriteString(fmt.Sprintf("  %s\n\n", installCmd))
		
		message.WriteString("Or use compile mode to generate assembly only:\n")
		message.WriteString("  c65gm compile -i program.c65\n\n")
		
		message.WriteString("For manual download or more information:\n")
		message.WriteString("  https://github.com/meonwax/acme")
		
		return fmt.Errorf("%s", message.String())
	}
	return nil
}