package erf

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"regexp"
	"sort"
	"strings"
)

var expectedSHA256Pattern = regexp.MustCompile(`^[0-9a-f]{64}$`)

const (
	headerSize = 160
	versionV10 = "V1.0"
	strRefNone = 0xFFFFFFFF
	typeMOD    = "MOD "
	typeHAK    = "HAK "
)

type Archive struct {
	FileType  string
	Version   string
	Resources []Resource
}

type Resource struct {
	Name       string
	Type       uint16
	Data       []byte
	SourcePath string
	Size       int64
	// ExpectedSHA256, when set, is the lowercase hex SHA-256 the streamed
	// SourcePath payload must hash to. A mismatch fails the write.
	ExpectedSHA256 string
}

type header struct {
	FileType              [4]byte
	Version               [4]byte
	LanguageCount         uint32
	LocalizedStringSize   uint32
	EntryCount            uint32
	LocalizedStringOffset uint32
	KeyListOffset         uint32
	ResourceListOffset    uint32
	BuildYear             uint32
	BuildDay              uint32
	DescriptionStrRef     uint32
	Reserved              [116]byte
}

type keyEntry struct {
	ResRef       [16]byte
	ResourceID   uint32
	ResourceType uint16
	Unused       uint16
}

type resourceEntry struct {
	Offset uint32
	Size   uint32
}

var extensionTypes = map[string]uint16{
	"res": 0x0000,
	"bmp": 0x0001,
	"mve": 0x0002,
	"tga": 0x0003,
	"wav": 0x0004,
	"plt": 0x0006,
	"ini": 0x0007,
	"bmu": 0x0008,
	"txt": 0x000A,
	"mdl": 0x07D2,
	"nss": 0x07D9,
	"ncs": 0x07DA,
	"are": 0x07DC,
	"set": 0x07DD,
	"ifo": 0x07DE,
	"bic": 0x07DF,
	"wok": 0x07E0,
	"2da": 0x07E1,
	"tlk": 0x07E2,
	"txi": 0x07E6,
	"git": 0x07E7,
	"uti": 0x07E9,
	"utc": 0x07EB,
	"dlg": 0x07ED,
	"itp": 0x07EE,
	"utt": 0x07F0,
	"dds": 0x07F1,
	"uts": 0x07F3,
	"fac": 0x07F6,
	"gff": 0x07F7,
	"ute": 0x07F8,
	"utd": 0x07FA,
	"utp": 0x07FC,
	"dfa": 0x07FD,
	"gic": 0x07FE,
	"gui": 0x07FF,
	"utm": 0x0803,
	"dwk": 0x0804,
	"pwk": 0x0805,
	"utg": 0x0807,
	"jrl": 0x0808,
	"utw": 0x080A,
	"ssf": 0x080C,
	"hak": 0x080D,
	"nwm": 0x080E,
	"bik": 0x080F,
	"ndb": 0x0810,
	"ptm": 0x0811,
	"ptt": 0x0812,
	"ltr": 0x0813,
	"shd": 0x0815,
	"mdb": 0x0816,
	"mtr": 0x0818,
	"jpg": 0x081C,
	"lod": 0x081E,
	"png": 0x0820,
	"lyt": 0x0BB8,
	"vis": 0x0BB9,
	"mdx": 0x0BC0,
	"wlk": 0x0BCC,
	"xml": 0x0BCD,
	"gr2": 0x0FA3,
}

var typeExtensions = map[uint16]string{
	0x0000: "res",
	0x0001: "bmp",
	0x0002: "mve",
	0x0003: "tga",
	0x0004: "wav",
	0x0006: "plt",
	0x0007: "ini",
	0x0008: "bmu",
	0x000A: "txt",
	0x07D2: "mdl",
	0x07D9: "nss",
	0x07DA: "ncs",
	0x07DC: "are",
	0x07DD: "set",
	0x07DE: "ifo",
	0x07DF: "bic",
	0x07E0: "wok",
	0x07E1: "2da",
	0x07E2: "tlk",
	0x07E6: "txi",
	0x07E7: "git",
	0x07E9: "uti",
	0x07EB: "utc",
	0x07ED: "dlg",
	0x07EE: "itp",
	0x07F0: "utt",
	0x07F1: "dds",
	0x07F3: "uts",
	0x07F6: "fac",
	0x07F7: "gff",
	0x07F8: "ute",
	0x07FA: "utd",
	0x07FC: "utp",
	0x07FD: "dfa",
	0x07FE: "gic",
	0x07FF: "gui",
	0x0803: "utm",
	0x0804: "dwk",
	0x0805: "pwk",
	0x0807: "utg",
	0x0808: "jrl",
	0x080A: "utw",
	0x080C: "ssf",
	0x080D: "hak",
	0x080E: "nwm",
	0x080F: "bik",
	0x0810: "ndb",
	0x0811: "ptm",
	0x0812: "ptt",
	0x0813: "ltr",
	0x0815: "shd",
	0x0816: "mdb",
	0x0818: "mtr",
	0x081C: "jpg",
	0x081E: "lod",
	0x0820: "png",
	0x0BB8: "lyt",
	0x0BB9: "vis",
	0x0BC0: "mdx",
	0x0BCC: "wlk",
	0x0BCD: "xml",
	0x0FA3: "gr2",
}

func init() {
	for ext, resourceType := range extensionTypes {
		canonicalExt, ok := typeExtensions[resourceType]
		if !ok {
			panic(fmt.Sprintf("missing canonical extension for resource type 0x%04X", resourceType))
		}
		if canonicalExt == ext {
			continue
		}
	}
}

func New(fileType string, resources []Resource) Archive {
	normalized := make([]Resource, len(resources))
	copy(normalized, resources)
	sort.Slice(normalized, func(i, j int) bool {
		if normalized[i].Name == normalized[j].Name {
			return normalized[i].Type < normalized[j].Type
		}
		return normalized[i].Name < normalized[j].Name
	})

	return Archive{
		FileType:  padFour(fileType),
		Version:   versionV10,
		Resources: normalized,
	}
}

func ArchiveSize(resources []Resource) int64 {
	return int64(headerSize+len(resources)*24+len(resources)*8) + totalResourceBytes(resources)
}

func Write(w io.Writer, archive Archive) error {
	if archive.Version == "" {
		archive.Version = versionV10
	}
	if archive.FileType == "" {
		archive.FileType = typeMOD
	}

	keys := make([]keyEntry, 0, len(archive.Resources))
	entries := make([]resourceEntry, 0, len(archive.Resources))

	dataOffset := uint32(headerSize + len(archive.Resources)*24 + len(archive.Resources)*8)
	for index, resource := range archive.Resources {
		if len(resource.Name) > 16 {
			return fmt.Errorf("resource %q exceeds 16-byte resref limit", resource.Name)
		}
		size, err := payloadSize(resource)
		if err != nil {
			return err
		}
		if size > int64(^uint32(0)) {
			return fmt.Errorf("resource %q exceeds 4 GiB resource size limit", resource.Name)
		}

		entry := resourceEntry{
			Offset: dataOffset,
			Size:   uint32(size),
		}
		entries = append(entries, entry)
		dataOffset += uint32(size)

		var key keyEntry
		copy(key.ResRef[:], []byte(resource.Name))
		key.ResourceID = uint32(index)
		key.ResourceType = resource.Type
		keys = append(keys, key)
	}

	hdr := header{
		LanguageCount:         0,
		LocalizedStringSize:   0,
		EntryCount:            uint32(len(archive.Resources)),
		LocalizedStringOffset: headerSize,
		KeyListOffset:         headerSize,
		ResourceListOffset:    uint32(headerSize + len(keys)*24),
		BuildYear:             0,
		BuildDay:              0,
		DescriptionStrRef:     strRefNone,
	}
	copy(hdr.FileType[:], []byte(padFour(archive.FileType)))
	copy(hdr.Version[:], []byte(padFour(archive.Version)))

	if err := binary.Write(w, binary.LittleEndian, hdr); err != nil {
		return err
	}
	if err := binary.Write(w, binary.LittleEndian, keys); err != nil {
		return err
	}
	if err := binary.Write(w, binary.LittleEndian, entries); err != nil {
		return err
	}
	for _, resource := range archive.Resources {
		if err := writeResourceData(w, resource); err != nil {
			return err
		}
	}
	return nil
}

func Read(r io.Reader) (Archive, error) {
	data, err := io.ReadAll(r)
	if err != nil {
		return Archive{}, fmt.Errorf("read erf: %w", err)
	}
	if len(data) < headerSize {
		return Archive{}, fmt.Errorf("erf file too small: %d bytes", len(data))
	}

	var hdr header
	if err := binary.Read(bytes.NewReader(data[:headerSize]), binary.LittleEndian, &hdr); err != nil {
		return Archive{}, fmt.Errorf("decode erf header: %w", err)
	}

	keyStart := int(hdr.KeyListOffset)
	keyEnd := keyStart + int(hdr.EntryCount)*24
	if keyEnd > len(data) {
		return Archive{}, fmt.Errorf("erf key list exceeds file bounds")
	}
	keys := make([]keyEntry, hdr.EntryCount)
	if err := binary.Read(bytes.NewReader(data[keyStart:keyEnd]), binary.LittleEndian, &keys); err != nil {
		return Archive{}, fmt.Errorf("decode key list: %w", err)
	}

	resourceStart := int(hdr.ResourceListOffset)
	resourceEnd := resourceStart + int(hdr.EntryCount)*8
	if resourceEnd > len(data) {
		return Archive{}, fmt.Errorf("erf resource list exceeds file bounds")
	}
	entries := make([]resourceEntry, hdr.EntryCount)
	if err := binary.Read(bytes.NewReader(data[resourceStart:resourceEnd]), binary.LittleEndian, &entries); err != nil {
		return Archive{}, fmt.Errorf("decode resource list: %w", err)
	}

	resources := make([]Resource, 0, hdr.EntryCount)
	for index, key := range keys {
		entry := entries[index]
		start := int(entry.Offset)
		end := start + int(entry.Size)
		if end > len(data) {
			return Archive{}, fmt.Errorf("resource %d exceeds file bounds", index)
		}

		resref := string(bytes.TrimRight(key.ResRef[:], "\x00"))
		payload := make([]byte, entry.Size)
		copy(payload, data[start:end])
		resources = append(resources, Resource{
			Name: resref,
			Type: key.ResourceType,
			Data: payload,
			Size: int64(entry.Size),
		})
	}

	return Archive{
		FileType:  string(hdr.FileType[:]),
		Version:   string(hdr.Version[:]),
		Resources: resources,
	}, nil
}

func ResourceTypeForExtension(extension string) (uint16, bool) {
	resourceType, ok := extensionTypes[strings.TrimPrefix(strings.ToLower(extension), ".")]
	return resourceType, ok
}

func HAKResourceTypeForExtension(extension string) (uint16, bool) {
	return ResourceTypeForExtension(extension)
}

func ExtensionForResourceType(resourceType uint16) (string, bool) {
	extension, ok := typeExtensions[resourceType]
	return extension, ok
}

func padFour(value string) string {
	value = strings.ToUpper(value)
	if len(value) >= 4 {
		return value[:4]
	}
	return value + strings.Repeat(" ", 4-len(value))
}

func totalResourceBytes(resources []Resource) int64 {
	var total int64
	for _, resource := range resources {
		switch {
		case resource.Size > 0:
			total += resource.Size
		default:
			total += int64(len(resource.Data))
		}
	}
	return total
}

func payloadSize(resource Resource) (int64, error) {
	if resource.Size > 0 {
		return resource.Size, nil
	}
	if resource.SourcePath != "" && len(resource.Data) == 0 {
		info, err := os.Stat(resource.SourcePath)
		if err != nil {
			return 0, fmt.Errorf("stat resource %q: %w", resource.SourcePath, err)
		}
		return info.Size(), nil
	}
	return int64(len(resource.Data)), nil
}

func writeResourceData(w io.Writer, resource Resource) error {
	if len(resource.Data) > 0 || resource.SourcePath == "" {
		_, err := w.Write(resource.Data)
		return err
	}

	if resource.ExpectedSHA256 != "" && !expectedSHA256Pattern.MatchString(resource.ExpectedSHA256) {
		return fmt.Errorf("resource %q has invalid expected sha256", resource.SourcePath)
	}

	file, err := os.Open(resource.SourcePath)
	if err != nil {
		return fmt.Errorf("open resource %q: %w", resource.SourcePath, err)
	}
	defer file.Close()

	hash := sha256.New()
	written, err := io.Copy(io.MultiWriter(w, hash), file)
	if err != nil {
		return fmt.Errorf("copy resource %q: %w", resource.SourcePath, err)
	}
	if resource.Size > 0 && written != resource.Size {
		return fmt.Errorf("copy resource %q: expected %d bytes, wrote %d", resource.SourcePath, resource.Size, written)
	}
	if resource.ExpectedSHA256 != "" {
		got := hex.EncodeToString(hash.Sum(nil))
		if got != resource.ExpectedSHA256 {
			return fmt.Errorf("resource %q sha256 mismatch: expected %s, got %s", resource.SourcePath, resource.ExpectedSHA256, got)
		}
	}
	return nil
}