package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"strconv"
)

type section struct {
	name string
	buff bytes.Buffer
}

type symtabEntry struct {
	sectionName string
	kind        string
	offset      int64
	global      bool
}

type compiler struct {
	symtab         map[string]symtabEntry
	sections       []section
	currentSection *section
}

func NewCompiler() *compiler {
	return &compiler{
		symtab: map[string]symtabEntry{}, // starts out empty
	}
}

func (c *compiler) Must(b []byte) {
	n, err := c.currentSection.buff.Write(b)
	if err != nil {
		panic(err)
	}
	if n != len(b) {
		panic(io.ErrShortWrite)
	}
}

func (c *compiler) MustUint64(val uint64) {
	ret := make([]byte, 8)
	binary.LittleEndian.PutUint64(ret, val)
	c.Must(ret)
}

func (c *compiler) Compile(t Token) error {
	if c.currentSection == nil {
		// The only allowable token outside of a section is to start a new section
		if _, ok := t.(SectionToken); !ok {
			return fmt.Errorf("Need to start with a section token, got %#t", t)
		}
	}

	switch tok := t.(type) {
	case SectionToken:

		// Check if we are resuming an existing section
		for i, sec := range c.sections {
			if sec.name == tok.SectionName {
				// Found it
				c.currentSection = &c.sections[i]
				return nil
			}
		}

		// It's a new section
		c.sections = append(c.sections, section{
			name: tok.SectionName,
			buff: bytes.Buffer{},
		})
		c.currentSection = &c.sections[len(c.sections)-1]

		return nil

	case DataVariableInstrToken:
		// Stash in symbol table for future backreferences
		if _, ok := c.symtab[tok.VarName]; ok {
			return fmt.Errorf("variable %q was already defined", tok.VarName)
		}
		c.symtab[tok.VarName] = symtabEntry{
			sectionName: c.currentSection.name,
			kind:        ".var." + tok.Sizeclass,
			offset:      int64(c.currentSection.buff.Len()),
			global:      false, // TODO allow this?
		}

		// Generate bytes for the symbol
		switch tok.Sizeclass {
		case "u8":
			// 1 byte literal
			val, err := strconv.ParseUint(tok.Value, 10, 8)
			if err != nil {
				return err
			}

			c.Must([]byte{byte(val)})
			return nil

		case "u64":
			// 8-byte literal
			val, err := strconv.ParseUint(tok.Value, 10, 64)
			if err != nil {
				return err
			}

			c.MustUint64(val)
			return nil

		case "sz":
			// string with null termination
			ret := []byte(tok.Value)
			ret = append(ret, 0)
			c.Must(ret)
			return nil

		default:
			return fmt.Errorf("variable %q has unknown size class %q", tok.VarName, tok.Sizeclass)
		}

	case LabelToken:
		if _, ok := c.symtab[tok.LabelName]; ok {
			return fmt.Errorf("name %q was already defined", tok.LabelName)
		}
		c.symtab[tok.LabelName] = symtabEntry{
			sectionName: c.currentSection.name,
			kind:        ".label",
			offset:      int64(c.currentSection.buff.Len()),
			global:      tok.IsGlobal,
		}
		return nil

	case MovInstrToken:
		// TODO encode more cases properly
		if literal, err := strconv.ParseInt(tok.Args[1], 10, 64); err == nil {

			// Store immediate in register
			switch tok.Args[0] {
			case "rax":
				c.Must([]byte{0x48, 0xb8}) // TODO store in eax with shorter prefix if <32 bit
				c.MustUint64(uint64(literal))

			case "rbx":
				c.Must([]byte{0x48, 0xbb}) // TODO store in eax with shorter prefix if <32 bit
				c.MustUint64(uint64(literal))

			case "rcx":
				c.Must([]byte{0x48, 0xb9}) // TODO store in eax with shorter prefix if <32 bit
				c.MustUint64(uint64(literal))

			case "rdx":
				c.Must([]byte{0x48, 0xba}) // TODO store in eax with shorter prefix if <32 bit
				c.MustUint64(uint64(literal))

			default:
				// Store immediate in variable?
				panic("not implemented: store immediate in ???? thing")
			}

		} else if _, ok := c.symtab[tok.Args[1]]; ok {
			// Store variable's contents in register

		} else if _, ok := c.symtab["&"+tok.Args[1]]; ok {
			// With &; store address of variable in register

		}

		panic("unknown mov type, sorry")

	default:
		return fmt.Errorf("can't compile token of type %#t", t)
	}
}

func (c *compiler) Finalize(dest io.Writer) error {

	const alignment = 4096

	// Write ELF header
	// Write section headers
	// Write binary content
	// Pad out section to page alignment
	// Done

	panic("TODO")
}