package commands

import (
	"fmt"
	"os"
	"strings"

	"c65gm/internal/compiler"
	"c65gm/internal/preproc"
	"c65gm/internal/utils"
)

// ShiftRCommand handles logical shift right operations
// Syntax:
//
//	SHIFTR <source> BY <amount> GIVING <dest>    # old syntax with BY/GIVING
//	SHIFTR <source> >> <amount> -> <dest>        # old syntax with >>/->
//	<dest> = <source> >> <amount>                # new syntax
type ShiftRCommand struct {
	sourceVarName string
	sourceVarKind compiler.VarKind
	sourceValue   uint16
	sourceIsVar   bool

	amountVarName string
	amountVarKind compiler.VarKind
	amountValue   uint16
	amountIsVar   bool

	destVarName string
	destVarKind compiler.VarKind

	line preproc.Line // Store line info for warnings
}

func (c *ShiftRCommand) WillHandle(line preproc.Line) bool {
	params, err := utils.ParseParams(line.Text)
	if err != nil || len(params) == 0 {
		return false
	}

	// Old syntax: SHIFTR ... (must have exactly 6 params)
	if strings.ToUpper(params[0]) == "SHIFTR" && len(params) == 6 {
		return true
	}

	// New syntax: <dest> = <source> >> <amount>
	if len(params) == 5 && params[1] == "=" && params[3] == ">>" {
		return true
	}

	return false
}

func (c *ShiftRCommand) Interpret(line preproc.Line, ctx *compiler.CompilerContext) error {
	// Clear state
	c.sourceVarName = ""
	c.sourceIsVar = false
	c.sourceValue = 0
	c.amountVarName = ""
	c.amountIsVar = false
	c.amountValue = 0
	c.destVarName = ""
	c.line = line // Store line for warnings

	params, err := utils.ParseParams(line.Text)
	if err != nil {
		return err
	}

	paramCount := len(params)
	scope := ctx.CurrentScope()

	// Create constant lookup function
	constLookup := ctx.SymbolTable.ConstantLookupFunc(scope)

	// Determine syntax and parse accordingly
	if strings.ToUpper(params[0]) == "SHIFTR" {
		// Old syntax: SHIFTR <source> BY/>> <amount> GIVING/-> <dest>
		if paramCount != 6 {
			return fmt.Errorf("SHIFTR: wrong number of parameters (%d), expected 6", paramCount)
		}

		separator1 := strings.ToUpper(params[2])
		if separator1 != "BY" && separator1 != ">>" {
			return fmt.Errorf("SHIFTR: parameter #3 must be 'BY' or '>>', got %q", params[2])
		}

		separator2 := strings.ToUpper(params[4])
		if separator2 != "GIVING" && separator2 != "->" {
			return fmt.Errorf("SHIFTR: parameter #5 must be 'GIVING' or '->', got %q", params[4])
		}

		// Parse destination
		destName := params[5]
		destSym := ctx.SymbolTable.Lookup(destName, scope)
		if destSym == nil {
			return fmt.Errorf("SHIFTR: unknown variable %q", destName)
		}
		if destSym.IsConst() {
			return fmt.Errorf("SHIFTR: cannot assign to constant %q", destName)
		}
		c.destVarName = destSym.FullName()
		c.destVarKind = destSym.GetVarKind()

		// Parse source
		var err error
		c.sourceVarName, c.sourceVarKind, c.sourceValue, c.sourceIsVar, err = compiler.ParseOperandParam(
			params[1], ctx.SymbolTable, scope, constLookup)
		if err != nil {
			return fmt.Errorf("SHIFTR: source: %w", err)
		}

		// Parse amount
		c.amountVarName, c.amountVarKind, c.amountValue, c.amountIsVar, err = compiler.ParseOperandParam(
			params[3], ctx.SymbolTable, scope, constLookup)
		if err != nil {
			return fmt.Errorf("SHIFTR: amount: %w", err)
		}

	} else {
		// New syntax: <dest> = <source> >> <amount>
		if paramCount != 5 {
			return fmt.Errorf("SHIFTR: wrong number of parameters (%d), expected 5", paramCount)
		}

		if params[1] != "=" {
			return fmt.Errorf("SHIFTR: expected '=' at position 2, got %q", params[1])
		}

		if params[3] != ">>" {
			return fmt.Errorf("SHIFTR: expected '>>' at position 4, got %q", params[3])
		}

		// Parse destination
		destName := params[0]
		destSym := ctx.SymbolTable.Lookup(destName, scope)
		if destSym == nil {
			return fmt.Errorf("SHIFTR: unknown variable %q", destName)
		}
		if destSym.IsConst() {
			return fmt.Errorf("SHIFTR: cannot assign to constant %q", destName)
		}
		c.destVarName = destSym.FullName()
		c.destVarKind = destSym.GetVarKind()

		// Parse source
		var err error
		c.sourceVarName, c.sourceVarKind, c.sourceValue, c.sourceIsVar, err = compiler.ParseOperandParam(
			params[2], ctx.SymbolTable, scope, constLookup)
		if err != nil {
			return fmt.Errorf("SHIFTR: source: %w", err)
		}

		// Parse amount
		c.amountVarName, c.amountVarKind, c.amountValue, c.amountIsVar, err = compiler.ParseOperandParam(
			params[4], ctx.SymbolTable, scope, constLookup)
		if err != nil {
			return fmt.Errorf("SHIFTR: amount: %w", err)
		}
	}

	// Validate amount
	if c.amountIsVar {
		if c.amountVarKind == compiler.KindWord {
			return fmt.Errorf("SHIFTR: amount must be BYTE variable, got WORD %q", c.amountVarName)
		}
	} else {
		if c.amountValue > 255 {
			return fmt.Errorf("SHIFTR: amount constant %d out of BYTE range (0-255)", c.amountValue)
		}
	}

	return nil
}

func (c *ShiftRCommand) Generate(ctx *compiler.CompilerContext) ([]string, error) {
	// Case 1: BYTE destination
	if c.destVarKind == compiler.KindByte {
		return c.generateByteShift(ctx)
	}
	
	// Case 2: WORD destination
	return c.generateWordShift(ctx)
}

// generateByteShift handles BYTE >> amount
func (c *ShiftRCommand) generateByteShift(ctx *compiler.CompilerContext) ([]string, error) {
	// Special case: WORD source, BYTE destination
	if c.sourceVarKind == compiler.KindWord {
		return c.generateWordToByteShift(ctx)
	}
	
	var asm []string
	
	// Constant shift amount
	if !c.amountIsVar {
		amount := c.amountValue
		
		// Shift >= 8 bits -> zero
		if amount >= 8 {
			if amount > 0 {
				_, _ = fmt.Fprintf(os.Stderr, "%s:%d: warning: shift amount %d >= 8 bits, value will be zero\n",
					c.line.Filename, c.line.LineNo, amount)
			}
			
			// Store zero
			if c.sourceIsVar && c.sourceVarName == c.destVarName {
				// Same variable: just zero it
				asm = append(asm, "\tlda #0")
				asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
			} else {
				// Different: load zero and store
				asm = append(asm, "\tlda #0")
				asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
			}
			return asm, nil
		}
		
		// Shift 0 -> copy source
		if amount == 0 {
			return c.generateByteCopy(), nil
		}
		
		// Constant shift 1-7
		return c.generateByteShiftConst(amount), nil
	}
	
	// Variable shift amount
	return c.generateByteShiftVar(ctx)
}

// generateByteCopy copies source to destination (BYTE)
func (c *ShiftRCommand) generateByteCopy() []string {
	var asm []string
	
	if c.sourceIsVar {
		if c.sourceVarName == c.destVarName {
			// Same variable, no copy needed
			return asm
		}
		// Different variable
		if c.sourceVarKind == compiler.KindWord {
			// WORD source, BYTE destination - just take low byte
			asm = append(asm, fmt.Sprintf("\tlda %s", c.sourceVarName))
			asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
		} else {
			// BYTE source
			asm = append(asm, fmt.Sprintf("\tlda %s", c.sourceVarName))
			asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
		}
	} else {
		// Literal
		val := uint8(c.sourceValue & 0xFF)
		asm = append(asm, fmt.Sprintf("\tlda #$%02x", val))
		asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
	}
	
	return asm
}

// generateByteShiftConst generates BYTE >> constant (1-7)
func (c *ShiftRCommand) generateByteShiftConst(amount uint16) []string {
	var asm []string
	
	// Copy source to destination if needed
	copyAsm := c.generateByteCopy()
	asm = append(asm, copyAsm...)
	
	// Apply shift (right shift uses LSR)
	for i := uint16(0); i < amount; i++ {
		asm = append(asm, fmt.Sprintf("\tlsr %s", c.destVarName))
	}
	
	return asm
}

// generateByteShiftVar generates BYTE >> variable
func (c *ShiftRCommand) generateByteShiftVar(ctx *compiler.CompilerContext) ([]string, error) {
	var asm []string
	
	// Copy source to destination if needed
	copyAsm := c.generateByteCopy()
	asm = append(asm, copyAsm...)
	
	// Generate labels
	loopLabel := ctx.GeneralStack.Push()
	ctx.GeneralStack.Pop()
	doneLabel := ctx.GeneralStack.Push()
	ctx.GeneralStack.Pop()
	
	// Variable shift loop (right shift uses LSR)
	asm = append(asm, fmt.Sprintf("\tldx %s", c.amountVarName))
	asm = append(asm, fmt.Sprintf("\tbeq %s", doneLabel))
	asm = append(asm, loopLabel)
	asm = append(asm, fmt.Sprintf("\tlsr %s", c.destVarName))
	asm = append(asm, "\tdex")
	asm = append(asm, fmt.Sprintf("\tbne %s", loopLabel))
	asm = append(asm, doneLabel)
	
	return asm, nil
}

// generateWordToByteShift handles WORD >> amount -> BYTE
func (c *ShiftRCommand) generateWordToByteShift(ctx *compiler.CompilerContext) ([]string, error) {
	var asm []string
	
	// Constant shift amount
	if !c.amountIsVar {
		amount := c.amountValue
		
		// Shift >= 16 bits -> zero
		if amount >= 16 {
			if amount > 0 {
				_, _ = fmt.Fprintf(os.Stderr, "%s:%d: warning: shift amount %d >= 16 bits, value will be zero\n",
					c.line.Filename, c.line.LineNo, amount)
			}
			
			// Store zero
			asm = append(asm, "\tlda #0")
			asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
			return asm, nil
		}
		
		// Shift 0 -> just take low byte
		if amount == 0 {
			if c.sourceIsVar {
				asm = append(asm, fmt.Sprintf("\tlda %s", c.sourceVarName))
				asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
			} else {
				val := uint8(c.sourceValue & 0xFF)
				asm = append(asm, fmt.Sprintf("\tlda #$%02x", val))
				asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
			}
			return asm, nil
		}
		
		// For WORD -> BYTE right shift, we need to handle the full WORD shift
		// then take the low byte. This is more complex than BYTE shift.
		
		// Shift 8-15: result depends only on high byte
		if amount >= 8 && amount < 16 {
			remaining := amount - 8
			
			// Get high byte
			if c.sourceIsVar {
				asm = append(asm, fmt.Sprintf("\tlda %s+1", c.sourceVarName))
			} else {
				hi := uint8((c.sourceValue >> 8) & 0xFF)
				asm = append(asm, fmt.Sprintf("\tlda #$%02x", hi))
			}
			
			// Store to destination
			asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
			
			// Shift remaining bits
			for i := uint16(0); i < remaining; i++ {
				asm = append(asm, fmt.Sprintf("\tlsr %s", c.destVarName))
			}
			
			return asm, nil
		}
		
		// Shift 1-7: need full WORD shift, then take low byte
		// We'll use the destination as a temporary for the low byte
		// and handle high byte in A register
		
		if c.sourceIsVar {
			// Load low byte to destination
			asm = append(asm, fmt.Sprintf("\tlda %s", c.sourceVarName))
			asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
			
			// Load high byte to A
			asm = append(asm, fmt.Sprintf("\tlda %s+1", c.sourceVarName))
		} else {
			// Literal source
			lo := uint8(c.sourceValue & 0xFF)
			hi := uint8((c.sourceValue >> 8) & 0xFF)
			
			asm = append(asm, fmt.Sprintf("\tlda #$%02x", lo))
			asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
			asm = append(asm, fmt.Sprintf("\tlda #$%02x", hi))
		}
		
		// Apply shifts: lsr A (high byte), ror destination (low byte)
		for i := uint16(0); i < amount; i++ {
			asm = append(asm, "\tlsr")
			asm = append(asm, fmt.Sprintf("\tror %s", c.destVarName))
		}
		
		return asm, nil
	}
	
	// Variable shift amount - use general approach with loop
	// We need X for shift count, A for high byte, destination for low byte
	
	// Generate labels
	loopLabel := ctx.GeneralStack.Push()
	ctx.GeneralStack.Pop()
	doneLabel := ctx.GeneralStack.Push()
	ctx.GeneralStack.Pop()
	
	if c.sourceIsVar {
		// Load low byte to destination
		asm = append(asm, fmt.Sprintf("\tlda %s", c.sourceVarName))
		asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
		
		// Load high byte to A
		asm = append(asm, fmt.Sprintf("\tlda %s+1", c.sourceVarName))
	} else {
		// Literal source
		lo := uint8(c.sourceValue & 0xFF)
		hi := uint8((c.sourceValue >> 8) & 0xFF)
		
		asm = append(asm, fmt.Sprintf("\tlda #$%02x", lo))
		asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
		asm = append(asm, fmt.Sprintf("\tlda #$%02x", hi))
	}
	
	// Load shift amount to X
	asm = append(asm, fmt.Sprintf("\tldx %s", c.amountVarName))
	asm = append(asm, fmt.Sprintf("\tbeq %s", doneLabel))
	
	// Shift loop
	asm = append(asm, loopLabel)
	asm = append(asm, "\tlsr")
	asm = append(asm, fmt.Sprintf("\tror %s", c.destVarName))
	asm = append(asm, "\tdex")
	asm = append(asm, fmt.Sprintf("\tbne %s", loopLabel))
	asm = append(asm, doneLabel)
	
	return asm, nil
}

// generateWordShift handles WORD >> amount
func (c *ShiftRCommand) generateWordShift(ctx *compiler.CompilerContext) ([]string, error) {
	// Constant shift amount
	if !c.amountIsVar {
		amount := c.amountValue
		
		// Shift >= 16 bits -> zero
		if amount >= 16 {
			if amount > 0 {
				_, _ = fmt.Fprintf(os.Stderr, "%s:%d: warning: shift amount %d >= 16 bits, value will be zero\n",
					c.line.Filename, c.line.LineNo, amount)
			}
			
			// Store zero
			if c.sourceIsVar && c.sourceVarName == c.destVarName {
				// Same variable: zero both bytes
				asm := []string{
					"\tlda #0",
					fmt.Sprintf("\tsta %s", c.destVarName),
					fmt.Sprintf("\tsta %s+1", c.destVarName),
				}
				return asm, nil
			} else {
				// Different: load zero and store to both bytes
				asm := []string{
					"\tlda #0",
					fmt.Sprintf("\tsta %s", c.destVarName),
					fmt.Sprintf("\tsta %s+1", c.destVarName),
				}
				return asm, nil
			}
		}
		
		// Shift 0 -> copy source
		if amount == 0 {
			return c.generateWordCopy(), nil
		}
		
		// Constant shift 1-15
		return c.generateWordShiftConst(amount), nil
	}
	
	// Variable shift amount
	return c.generateWordShiftVar(ctx)
}

// generateWordCopy copies source to destination (WORD)
func (c *ShiftRCommand) generateWordCopy() []string {
	var asm []string
	
	if c.sourceIsVar && c.sourceVarName == c.destVarName {
		// Same variable, no copy needed
		return asm
	}
	
	if c.sourceIsVar {
		// Variable source
		if c.sourceVarKind == compiler.KindByte {
			// Byte -> Word (zero-extend)
			asm = append(asm, fmt.Sprintf("\tlda %s", c.sourceVarName))
			asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
			asm = append(asm, "\tlda #0")
			asm = append(asm, fmt.Sprintf("\tsta %s+1", c.destVarName))
		} else {
			// Word -> Word
			asm = append(asm, fmt.Sprintf("\tlda %s", c.sourceVarName))
			asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
			asm = append(asm, fmt.Sprintf("\tlda %s+1", c.sourceVarName))
			asm = append(asm, fmt.Sprintf("\tsta %s+1", c.destVarName))
		}
	} else {
		// Literal source
		lo := uint8(c.sourceValue & 0xFF)
		hi := uint8((c.sourceValue >> 8) & 0xFF)
		asm = append(asm, fmt.Sprintf("\tlda #$%02x", lo))
		asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
		if lo != hi {
			asm = append(asm, fmt.Sprintf("\tlda #$%02x", hi))
		}
		asm = append(asm, fmt.Sprintf("\tsta %s+1", c.destVarName))
	}
	
	return asm
}

// generateWordShiftConst generates WORD >> constant (1-15)
func (c *ShiftRCommand) generateWordShiftConst(amount uint16) []string {
	var asm []string
	
	// Special case: shift >= 8
	if amount >= 8 && amount < 16 {
		remaining := amount - 8
		
		// Get source high byte
		if c.sourceIsVar {
			if c.sourceVarName == c.destVarName {
				// Same variable: read from destination (already has value)
				asm = append(asm, fmt.Sprintf("\tlda %s+1", c.destVarName))
			} else {
				// Different: read from source
				asm = append(asm, fmt.Sprintf("\tlda %s+1", c.sourceVarName))
			}
		} else {
			// Literal
			hi := uint8((c.sourceValue >> 8) & 0xFF)
			asm = append(asm, fmt.Sprintf("\tlda #$%02x", hi))
		}
		
		// Store to destination low byte
		asm = append(asm, fmt.Sprintf("\tsta %s", c.destVarName))
		
		// Zero high byte
		asm = append(asm, "\tlda #0")
		asm = append(asm, fmt.Sprintf("\tsta %s+1", c.destVarName))
		
		// Shift low byte remaining bits (right shift)
		for i := uint16(0); i < remaining; i++ {
			asm = append(asm, fmt.Sprintf("\tlsr %s", c.destVarName))
		}
		
		return asm
	}
	
	// Shift 1-7: normal copy and shift
	copyAsm := c.generateWordCopy()
	asm = append(asm, copyAsm...)
	
	// Apply shift (right shift uses LSR high byte, ROR low byte)
	for i := uint16(0); i < amount; i++ {
		asm = append(asm, fmt.Sprintf("\tlsr %s+1", c.destVarName))
		asm = append(asm, fmt.Sprintf("\tror %s", c.destVarName))
	}
	
	return asm
}

// generateWordShiftVar generates WORD >> variable
func (c *ShiftRCommand) generateWordShiftVar(ctx *compiler.CompilerContext) ([]string, error) {
	var asm []string
	
	// Copy source to destination if needed
	copyAsm := c.generateWordCopy()
	asm = append(asm, copyAsm...)
	
	// Generate labels
	loopLabel := ctx.GeneralStack.Push()
	ctx.GeneralStack.Pop()
	doneLabel := ctx.GeneralStack.Push()
	ctx.GeneralStack.Pop()
	
	// Variable shift loop (right shift uses LSR high byte, ROR low byte)
	asm = append(asm, fmt.Sprintf("\tldx %s", c.amountVarName))
	asm = append(asm, fmt.Sprintf("\tbeq %s", doneLabel))
	asm = append(asm, loopLabel)
	asm = append(asm, fmt.Sprintf("\tlsr %s+1", c.destVarName))
	asm = append(asm, fmt.Sprintf("\tror %s", c.destVarName))
	asm = append(asm, "\tdex")
	asm = append(asm, fmt.Sprintf("\tbne %s", loopLabel))
	asm = append(asm, doneLabel)
	
	return asm, nil
}