package main

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

type section struct {
	name                string
	name_shstrtabOffset int
	buff                bytes.Buffer
}

type symtabEntry struct {
	symtabSectionIndex int

	sectionName string
	kind        string
	offset      int64
	length      int64
	global      bool
}

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

func NewCompiler() *compiler {
	c := &compiler{
		symtab: map[string]symtabEntry{},
	}

	c.sections = append(c.sections, section{
		name:                `.shstrtab`, // Mandatory: the table that names sections themselves
		name_shstrtabOffset: 1,
		buff:                bytes.Buffer{},
	})
	c.shstrtab = &c.sections[0]
	// The first byte in a string table is conventionally expected be \x00, so that you can reference
	// null strings with it
	c.shstrtab.buff.WriteByte(0)

	c.shstrtab.buff.WriteString(c.shstrtab.name)
	c.shstrtab.buff.WriteByte(0)

	return c
}

func (c *compiler) CreateSymbol(name string, class string, offset int64, length int64, global bool) error {

	if _, ok := c.symtab[name]; ok {
		return fmt.Errorf("Symbol %q already exists", name)
	}

	// Find the .symtab section, or create if it does not exist
	symtabSec := c.FindOrCreateSection(`.symtab`)

	// New entry index = length / len(entry) = length / 24
	nextIndex := symtabSec.buff.Len() / 24

	// Add to our fast lookup table
	c.symtab[name] = symtabEntry{
		symtabSectionIndex: nextIndex,

		sectionName: c.currentSection.name,
		kind:        class,
		offset:      offset,
		global:      global,
		length:      length,
	}

	// Find the section index for the section containing this symbol
	sectionIndex := -1
	for i, _ := range c.sections {
		if c.sections[i].name == c.currentSection.name {
			sectionIndex = i
			break
		}
	}
	if sectionIndex == -1 {
		return fmt.Errorf("Current section missing index")
	}

	// Add to the .symtab section
	// This is required for variable references - after our single ELF .o is
	// created, linking it with any other .o files will create a combined .text
	// section where all the offsets have shifted
	esym := Elf64_Sym{}
	esym.st_name = 0 // Default: unnamed
	esym.st_info = STT_SECTION | (STB_LOCAL << 4)
	esym.st_other = STV_HIDDEN // For this translation unit only
	esym.st_shndx = uint16(sectionIndex)
	esym.st_size = uint64(length)

	err := binary.Write(&symtabSec.buff, binary.LittleEndian, &esym)
	return err
}

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

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

func (c *compiler) FindOrCreateSection(sectionName string) *section {

	if len(sectionName) == 0 || sectionName[0] != '.' {
		panic("section name should start with leading period")
	}

	for i, sec := range c.sections {
		if sec.name != sectionName {
			continue
		}

		// found it
		return &c.sections[i]
	}

	// No section with this name. Create it
	c.sections = append(c.sections, section{
		name:                sectionName,
		name_shstrtabOffset: c.shstrtab.buff.Len(),
		buff:                bytes.Buffer{},
	})

	c.shstrtab.buff.WriteString(sectionName)
	c.shstrtab.buff.WriteByte(0)

	return &c.sections[len(c.sections)-1]
}

func (c *compiler) Reloc(symbolName string, mode ElfRelocationType, addOffset int64) error {
	// Find '.rela.{currentsection}', creating it if it does not exist
	var relaSec *section = c.FindOrCreateSection(`.rela` + c.currentSection.name)

	// Find target symbol
	syminfo, ok := c.symtab[symbolName]
	if !ok {
		return fmt.Errorf("Reference to unknown symbol %q", symbolName)
	}

	// Add the relocation to the .rela section
	rr := Elf64_Rela{}
	rr.r_offset = uint64(c.currentSection.buff.Len())
	rr.r_info = uint64(syminfo.symtabSectionIndex)<<32 | uint64(mode) // high bits: Index of search symbol in the symtab. low bits: mode type
	rr.r_addend = addOffset

	err := binary.Write(&relaSec.buff, binary.LittleEndian, &rr)
	if err != nil {
		return err
	}

	// Done
	return nil
}

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:
		c.currentSection = c.FindOrCreateSection(tok.SectionName)
		return nil

	case DataVariableInstrToken:
		// Stash in symbol table for future backreferences
		// TODO allow making global symbols?
		// CreateSymbol does check for duplicate names already

		position := int64(c.currentSection.buff.Len())

		// 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)})

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

			c.MustUint64(val)

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

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

		err := c.CreateSymbol(tok.VarName, ".var."+tok.Sizeclass, int64(c.currentSection.buff.Len()), position-int64(c.currentSection.buff.Len()), false)
		if err != nil {
			return err
		}
		return nil

	case LabelToken:
		return c.CreateSymbol(tok.LabelName, ".label", int64(c.currentSection.buff.Len()), 0, tok.IsGlobal)

	case MovInstrToken:
		// TODO encode more cases properly
		if literal, err := strconv.ParseInt(tok.Args[1], 10, 64); err == nil {
			// mov rxx, imm
			// 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))

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

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

			default:
				// Store immediate in variable?
				panic("mov rxx,imm pattern: missing case")
			}
			return nil

		} else if strings.HasPrefix(tok.Args[0], `$`) {
			// mov $var, rxx
			// Load register's contents into variable
			// x86_64 can only really do this in a single instruction with 32-bit displacement, not full 64-bit

			switch tok.Args[1] {
			case "rax":
				c.Must([]byte{0x48, 0x89, 0x04, 0x25})
			default:
				panic("mov $var,rax pattern: missing case")
			}

			err = c.Reloc(tok.Args[0][1:], R_X86_64_32S, 0) // Declare that this is a 32-bit reloc, not a 64-bit one
			if err != nil {
				return fmt.Errorf("mov with relocation: %w", err)
			}
			c.Must([]byte{0, 0, 0, 0}) // 32-bit
			return nil

		} else if strings.HasPrefix(tok.Args[1], `$`) {
			// mov rxx, $var
			// With $; load variable contents into register

			switch tok.Args[0] {
			case "rax":
				c.Must([]byte{0x48, 0x8b, 0x04, 0x25})
			case "rdi":
				c.Must([]byte{0x48, 0x8b, 0x3c, 0x25})
			default:
				panic("mov rxx,$var pattern: missing case")
			}

			err = c.Reloc(tok.Args[1][1:], R_X86_64_32S, 0) // Declare that this is a 32-bit reloc, not a 64-bit one
			if err != nil {
				return fmt.Errorf("mov with relocation: %w", err)
			}
			c.Must([]byte{0, 0, 0, 0}) // 32-bit

			return nil

		} else if strings.HasPrefix(tok.Args[1], `&$`) {
			// mov rxx, &$var
			// With &; assign exact address of variable to register
			// This creates a movabs literal & a relocation entry
			// It's always 64-bit

			switch tok.Args[0] {
			case "rax":
				c.Must([]byte{0x48, 0xb8}) // TODO store in eax with shorter prefix if <32 bit
			case "rsi":
				c.Must([]byte{0x48, 0xbe}) // TODO store in eax with shorter prefix if <32 bit
			case "rdi":
				c.Must([]byte{0x48, 0xbf}) // TODO store in eax with shorter prefix if <32 bit

			default:
				panic("mov $var,rxx pattern: missing case")
			}

			err = c.Reloc(tok.Args[1][2:], R_X86_64_64, 0)
			if err != nil {
				return fmt.Errorf("mov with relocation: %w", err)
			}

			c.MustUint64(0)
			return nil

		} else if strings.HasPrefix(tok.Args[1], `strlen($`) && strings.HasSuffix(tok.Args[1], `)`) {
			// mov rxx, strlen($var)
			// With strlen; if this is an sz symbol, supply its length
			symname := tok.Args[1][8 : len(tok.Args[1])-1]
			sym, ok := c.symtab[symname]
			if !ok {
				return fmt.Errorf("Can't strlen on unknown variable %q", symname)
			}

			if sym.kind != ".var.sz" {
				return fmt.Errorf("Can't take the strlen of variable %q with type %q (expected sz)", symname, sym.kind)
			}

			effective := sym.length
			return c.Compile(MovInstrToken{Args: []string{tok.Args[0], strconv.Itoa(int(effective))}})

		} else {
			panic("unknown mov type, sorry")

		}

	case SyscallInstrToken:
		c.Must([]byte{0x0f, 0x05}) // syscall
		return nil

	case RetInstrToken:
		c.Must([]byte{0xc3}) // ret
		return nil

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

// Finalize exports the compiled sections into an ELF artefact.
// The resulting ELF is not executable directly, but it can be once fully
// linked (adding a program header and page alignment)
func (c *compiler) Finalize(dest io.Writer) error {

	// Write ELF header
	ehdr := Elf64_Ehdr{}
	ehdr.e_ident[0] = 0x7f
	ehdr.e_ident[1] = 'E'
	ehdr.e_ident[2] = 'L'
	ehdr.e_ident[3] = 'F'
	ehdr.e_ident[4] = 2 // 64-bit format
	ehdr.e_ident[5] = 1 // little endian
	ehdr.e_ident[6] = 1 // ELFv1 is the only format
	ehdr.e_ident[7] = 3 // Linux-compatible ABI

	ehdr.e_type = 0       // ET_NONE
	ehdr.e_machine = 0x3E // x86_64
	ehdr.e_version = 1    // ELFv1 again

	ehdr.e_shoff = 64     // The Ehdr is 64 bytes long, sections start immediately following
	ehdr.e_shentsize = 64 // Each Shdr is also 64 bytes long
	ehdr.e_shnum = uint16(len(c.sections))
	ehdr.e_shstrndx = 0 // We always put the .shstrtab as the 0th section

	err := binary.Write(dest, binary.LittleEndian, &ehdr)
	if err != nil {
		return err
	}

	// Don't declare a program header

	// Write section headers
	pctr := 64 + (64 * len(c.sections))
	for _, sec := range c.sections {
		shdr := Elf64_Shdr{}

		shdr.sh_name = uint32(sec.name_shstrtabOffset)
		switch sec.name {
		case ".text":
			shdr.sh_type = 1                 // SHT_PROGBITS, program data
			shdr.sh_flags = 0x2 | 0x4 | 0x10 // WRITE|ALLOC|MERGE
		case ".data":
			shdr.sh_type = 1           // SHT_PROGBITS, program data
			shdr.sh_flags = 0x2 | 0x10 // WRITE|MERGE
		case ".symtab":
			shdr.sh_type = 2            // SHT_SYMTAB
			shdr.sh_flags = 0x10 | 0x20 // MERGE|STRINGS
		case ".shstrtab":
			shdr.sh_type = 3            // SHT_STRTAB
			shdr.sh_flags = 0x10 | 0x20 // MERGE|STRINGS
		case ".rodata":
			fallthrough
		default: // Treat anything unknown as read-only data
			shdr.sh_type = 1     // SHT_PROGBITS, program data
			shdr.sh_flags = 0x10 // MERGE
		}

		shdr.sh_offset = uint64(pctr)
		shdr.sh_size = uint64(sec.buff.Len())

		pctr += sec.buff.Len()

		err = binary.Write(dest, binary.LittleEndian, &shdr)
		if err != nil {
			return err
		}
	}

	// Write binary content
	for _, sec := range c.sections {
		expectLen := sec.buff.Len()
		n, err := sec.buff.WriteTo(dest)
		if err != nil {
			return err
		}
		if n != int64(expectLen) {
			return io.ErrShortWrite
		}
	}

	// Done
	return nil

}