package main

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"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.

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 if arg == "-o" || arg == "-out" {
			if i+1 < len(args) {
				outputFile = args[i+1]
				i++ // Skip next arg
			}
		} 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
	}
	
	// Create temporary assembly file
	tempDir := os.TempDir()
	asmFile := filepath.Join(tempDir, "c65gm_"+filepath.Base(inputFile)+".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)
	}
	
	// Clean up temporary assembly file unless requested to keep it
	if !keepAsm {
		os.Remove(asmFile)
	} 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. Please install ACME: https://github.com/meonwax/acme")
	}
	
	// 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
}

// checkACMEAvailable checks if ACME is installed and provides helpful instructions if not
func checkACMEAvailable() error {
	_, err := exec.LookPath("acme")
	if err != nil {
		return fmt.Errorf(`ACME assembler not found in PATH.

c65gm requires ACME (Cross-Assembler) to create executable files.
Please install ACME from: https://github.com/meonwax/acme

Installation options:
- Linux: Download from releases or use package manager
- macOS: brew install acme
- Windows: Download binary from releases

Or use compile mode to generate assembly only:
  c65gm compile -i program.c65`)
	}
	return nil
}