package bundle

import (
	"bytes"
	"encoding/asn1"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"slices"
	"sort"
	"strconv"
	"strings"

	"github.com/apex/log"
	"github.com/blacktop/go-macho/types"
)

const Magic = "BUND"

type Header struct {
	Unknown1 uint16  // 0x0200
	Unknown2 uint16  // 0x1400
	_        uint32  // padding ?
	Magic    [4]byte // "BUND"
	_        uint16  // padding ?
	Type     uint16  // 3 (AOP/DCP), 4 (ExclaveCore)
}

type Bundle struct {
	Header
	TypeHeader any
	Config     Config
	Files      []File

	r      io.ReadSeeker
	closer any
}

func (b Bundle) String() string {
	var s strings.Builder
	s.WriteString(fmt.Sprintf("Bundle: %s\n", string(b.Magic[:])))
	s.WriteString(fmt.Sprintf("  Type: %d\n", b.Type))
	switch b.Type {
	case 3:
		s.WriteString("  Config:\n")
		s.WriteString(fmt.Sprintf("    Unk1: %d\n", b.Config.Unk1))
		s.WriteString(fmt.Sprintf("    Unk2: %d\n", b.Config.Unk2))
		s.WriteString("    Assets:\n")
		for i, h := range b.Config.Assets {
			s.WriteString(fmt.Sprintf("      %3s) %-20s\n", fmt.Sprintf("%d", i+1), h))
		}
		s.WriteString("    TOC:\n")
		for _, t := range b.Config.TOC {
			s.WriteString(fmt.Sprintf("      %s\n", t))
		}
		s.WriteString("Compartments:\n")
		for _, f := range b.Files {
			s.WriteString(fmt.Sprintf("%s\n", f))
		}
	case 4:
		s.WriteString("  Ranges:\n")
		for i, f := range b.TypeHeader.(Type4).Ranges {
			s.WriteString(fmt.Sprintf("    %3s) %s\n", fmt.Sprintf("%d", i+1), f))
		}
	}
	return s.String()
}

type Segment struct {
	Name   string
	Offset uint64
	Size   uint64
}

type Section struct {
	Name   string
	Offset uint64
	Size   uint64
}

type File struct {
	Name      string
	Type      string
	Segments  []Segment
	Sections  []Section
	Endpoints []Endpoint
	Extra     map[string]any
}

func (f File) Segment(name string) *Segment {
	for _, seg := range f.Segments {
		if seg.Name == name {
			return &seg
		}
	}
	return nil
}

func (f File) String() string {
	var s strings.Builder
	s.WriteString(fmt.Sprintf("  %s (%s)\n", f.Name, f.Type))
	for _, seg := range f.Segments {
		if seg.Size == 0 {
			continue
		}
		if seg.Name == "HEADER" {
			s.WriteString(fmt.Sprintf("    sz=0x%08x off=0x%08x-0x%08x __%s\n", seg.Size, seg.Offset, seg.Offset+seg.Size, seg.Name))
		}
	}
	for _, sec := range f.Sections {
		if sec.Size == 0 {
			continue
		}
		s.WriteString(fmt.Sprintf("      sz=0x%08x off=0x%08x-0x%08x __%s\n", sec.Size, sec.Offset, sec.Offset+sec.Size, sec.Name))
	}
	if len(f.Endpoints) > 0 {
		s.WriteString("    endpoints:\n")
		for i, ep := range f.Endpoints {
			s.WriteString(fmt.Sprintf("      %3s) %s\n", fmt.Sprintf("%d", i+1), ep))
		}
	}
	if len(f.Extra) > 0 {
		s.WriteString("    extra:\n")
		for k, v := range f.Extra {
			s.WriteString(fmt.Sprintf("      %s: %v\n", k, v))
		}
	}
	return s.String()
}

type Type3 struct {
	UUID             types.UUID
	_                [4]uint64 // padding ?
	SubType          uint64
	_                uint64 // padding ?
	FileSz           uint64
	_                uint64 // padding ?
	EndOff           uint64
	DataOff          uint64 // roottask __DATA offset
	_                uint64 // padding ?
	KernelDataOffset uint64 // 2D0C000h
	_                uint64 // padding ?
	TextOff          uint64 // roottask __TEXT offset
	UnkSz            uint64
	TextOffAgain     uint64 // ???
	TextSz           uint64 // is x14000 again
	UnkSzMaybe       uint64 // 2CF8000h
	_                uint64 // padding ?
	UnkOffMaybe      uint64 // 5FC000h
	UnkNumOfImgs     uint64 // C0000000h ? maybe just 0xC == 11 ?
	What             uint64 // C0035404h
	UnkSz2           uint64 // 5FC000h
	Something        uint64 // 26FC000h
	KernelDataOff    uint64 // 2D0C000h
	SomeSize2        uint64 // 334000h
	_                uint64 // padding ?
	AnotherSz        uint64 // 30000h ?
	_                uint64 // padding ?
	_                uint64 // padding ?
	YouAgain         uint64 // 30000h ?
	UnkSize4         uint64 // 74000h ?
	UnkSize5         uint64 // 40000h ?
	HelpMe           uint64 // C05FC000h ?
	UnkSize6         uint64 // 264000h ?
	Nooooooo         uint64 // 14000h ???
	SweetBaaaby      uint64 // C0670000h ??
	ImDead           uint64 // 70000h ??
	_                uint64 // padding ?
	Aaaagain         uint64 // 70000h ??
	UnkOffset6       uint64 // 2BC000h ???
	UnkOffset7       uint64 // 2D0C000h ???
	UnkOffset8       uint64 // 334000h ???
	Aaaaaaaaaagin    uint64 // 70000h ???
	_                uint64 // padding ?
	Aaaaaaaaaagin2   uint64 // 70000h ???
	UnkOffset9       uint64 // 2BC000h ???
	UnkOffset10      uint64 // 32C000h ???
	FooterOffset     uint64
	FooterSz         uint64
}

type Type4 struct {
	Unk0       uint32         // 1
	Unk1       uint32         // 0xc == 11
	_          uint64         // padding ?
	Unk2       uint64         // F000h
	_          [4]uint64      // padding ?
	Unk3       uint64         // C000h
	_          [2]uint64      // padding ?
	Unk4       uint64         // 3 ?
	_          uint64         // padding ?
	Unk5       uint64         // 16000h ?
	_          uint64         // padding ?
	_          uint64         // padding ?
	NumRanges  uint64         // 0xD == 12 ?
	UUID       types.UUID     // 636B62C3-4647-34F7-9089-A58256078A27
	_          uint64         // padding ?
	Unk7       uint64         // 14000h ?
	Unk7again  uint64         // 14000h ?
	Unk8       uint64         // 8000000h ?
	Unk9       uint64         // 1 ?
	_          uint64         // padding ?
	Unk10      uint64         // C000h ?
	Unk10again uint64         // C000h ?
	Unk11      uint64         // 8014000h ?
	Unk12      uint64         // 4 ?
	_          [3]uint64      // padding ?
	Unk13      uint64         // 8020000h ?
	Unk14      uint64         // 6 ?
	Unk15      uint64         // 8003E80h ?
	_          [36]uint64     // padding ?
	Unk17      uint64         // 0xa == 10 ?
	_          uint64         // padding ?
	Unk18      uint64         // 16000h ?
	Unk19      uint64         // 1582Ch ?
	_          uint64         // padding ?
	Unk20      uint64         // 0xa == 10 ?
	_          uint64         // padding ?
	Unk21      uint64         // F000h ?
	Ranges     [0xD]typ4Range // FIXME: this should be read AFTER the Type4 header is read
}

type typ4Range struct {
	Type   uint32
	Name   [4]byte
	Offset uint64
	Size   uint64
}

func (t4 typ4Range) String() string {
	slices.Reverse(t4.Name[:])
	if t4.Size == 0 {
		return fmt.Sprintf("typ=%d sz=%-10s off=0x%08x-0x%08x %s", t4.Type, fmt.Sprintf("%d", t4.Size), t4.Offset, t4.Offset+t4.Size, t4.Name)
	}
	return fmt.Sprintf("typ=%d sz=0x%08x off=0x%08x-0x%08x %s", t4.Type, t4.Size, t4.Offset, t4.Offset+t4.Size, t4.Name)
}

type Config struct {
	Unk1         int
	Unk2         int
	Assets       []Asset
	TOC          []TocEntry
	Compartments []Compartment
}

type Asset struct {
	Raw    asn1.RawContent
	Name   asn1.RawValue
	Type   int
	Offset int
	Size   int
}

func (h Asset) String() string {
	return fmt.Sprintf("%15s type=%d off=%#07x sz=%#x", h.Name.Bytes, h.Type, h.Offset, h.Size)
}

type TocEntry struct {
	Index int
	Entry asn1.RawValue `asn1:"optional"`
}

type TocEntryType struct {
	Name asn1.RawValue
	Type int
}

func (t TocEntry) GetEntry() *TocEntryType {
	if len(t.Entry.Bytes) > 0 {
		var typ TocEntryType
		if _, err := asn1.Unmarshal(t.Entry.Bytes, &typ); err == nil {
			return &typ
		}
	}
	return nil
}

func (t TocEntry) String() string {
	if entry := t.GetEntry(); entry != nil {
		return fmt.Sprintf("%3d) %15s type=%d", t.Index, entry.Name.Bytes, entry.Type)
	}
	return fmt.Sprintf("%3d) %s", t.Index, "nil")
}

type Compartment struct {
	Raw      asn1.RawContent
	AppUID   int
	Metadata []metadata
}

type metadata struct {
	Raw   asn1.RawContent
	Key   asn1.RawValue
	Value asn1.RawValue
}

type Endpoint struct {
	Type int
	Unk1 asn1.RawValue
	Unk2 asn1.RawValue
	Name asn1.RawValue
}

func (e Endpoint) String() string {
	return string(e.Name.Bytes)
}

func (md metadata) ParseValue() (any, error) {
	if bytes.HasPrefix(md.Key.Bytes, []byte("__COMPONENT")) {
		return string(md.Value.Bytes), nil
	}
	if bytes.HasPrefix(md.Key.Bytes, []byte("__ENDPOINT")) {
		var e Endpoint
		if _, err := asn1.Unmarshal(md.Value.Bytes, &e); err == nil {
			return e, nil
		} else {
			return nil, fmt.Errorf("failed to unmarshal bundle file info value: %v", err)
		}
	}
	if len(md.Value.Bytes) <= 8 {
		var num uint64
		for idx, b := range md.Value.Bytes {
			num |= uint64(b) << (8 * uint64(len(md.Value.Bytes)-1-idx))
		}
		return num, nil
	}
	return md.Value.Bytes, nil
}

func (md metadata) String() string {
	val, err := md.ParseValue()
	if err != nil {
		return fmt.Sprintf("[ERROR] failed to parse value: %v", err)
	}
	switch v := val.(type) {
	case string:
		return fmt.Sprintf("%s: %s", string(md.Key.Bytes), v)
	case Endpoint:
		return fmt.Sprintf("%s: %s", string(md.Key.Bytes), v)
	case uint64:
		if len(md.Value.Bytes) == 1 {
			return fmt.Sprintf("%s: %d", string(md.Key.Bytes), v)
		}
		return fmt.Sprintf("%s: %#x", string(md.Key.Bytes), v)
	default:
		return fmt.Sprintf("%s: %v", string(md.Key.Bytes), v)
	}
}

func (b *Bundle) ParseFiles() error {
	for _, bf := range b.Config.Compartments {
		var f File
		var sec Section
		var seg Segment
		entpoints := make(map[int]Endpoint, 0)
		f.Extra = make(map[string]any)
		for _, md := range bf.Metadata {
			val, err := md.ParseValue()
			if err != nil {
				return fmt.Errorf("failed to parse bundle file info value: %v", err)
			}
			if strings.EqualFold(string(md.Key.Bytes), "__COMPONENTNAME") {
				f.Name = val.(string)
			} else if strings.EqualFold(string(md.Key.Bytes), "__COMPONENTTYPE") {
				f.Type = val.(string)
			} else if _, secpart, ok := strings.Cut(string(md.Key.Bytes), "__MACHO__"); ok { // SECTION
				if name, _, ok := strings.Cut(secpart, "OFF"); ok {
					sec.Name = name
					sec.Offset = val.(uint64)
				}
				if name, _, ok := strings.Cut(secpart, "SZ"); ok {
					if sec.Name == "" {
						sec.Name = name
					}
					sec.Size = val.(uint64)
				}
				if sec.Name != "" && sec.Size != 0 {
					f.Sections = append(f.Sections, sec)
					sec = Section{}
				}
			} else if _, segpart, ok := strings.Cut(string(md.Key.Bytes), "__MACHO"); ok { // SEGMENT
				if name, _, ok := strings.Cut(segpart, "OFF"); ok {
					seg.Name = name
					seg.Offset = val.(uint64)
				}
				if name, _, ok := strings.Cut(segpart, "SZ"); ok {
					if seg.Name == "" {
						seg.Name = name
					}
					seg.Size = val.(uint64)
				}
				if seg.Name != "" && seg.Size != 0 {
					f.Segments = append(f.Segments, seg)
					seg = Segment{}
				}
			} else if _, idx, ok := strings.Cut(string(md.Key.Bytes), "__ENDPOINT__"); ok { // ENDPOINT
				if i, err := strconv.Atoi(idx); err == nil {
					entpoints[i] = val.(Endpoint)
				} else {
					return fmt.Errorf("failed to parse bundle file endpoint index: %v", err)
				}
			} else {
				f.Extra[string(md.Key.Bytes)] = string(md.Value.Bytes)
			}
		}
		sort.Slice(f.Sections, func(i, j int) bool {
			return f.Sections[i].Offset < f.Sections[j].Offset
		})
		sort.Slice(f.Segments, func(i, j int) bool {
			return f.Segments[i].Offset < f.Segments[j].Offset
		})
		keys := make([]int, 0, len(entpoints))
		for k := range entpoints {
			keys = append(keys, k)
		}
		sort.Ints(keys)
		for _, k := range keys {
			f.Endpoints = append(f.Endpoints, entpoints[k])
		}
		b.Files = append(b.Files, f)
	}
	return nil
}

func (b *Bundle) DumpFiles(output string) error {
	if b.Type == 4 {
		for _, rng := range b.TypeHeader.(Type4).Ranges {
			slices.Reverse(rng.Name[:])
			if rng.Size > 0 {
				if _, err := b.r.Seek(int64(rng.Offset), io.SeekStart); err != nil {
					return fmt.Errorf("failed to seek to bundle config data: %v", err)
				}
				data := make([]byte, rng.Size)
				if _, err := b.r.Read(data); err != nil {
					return fmt.Errorf("failed to read bundle data: %v", err)
				}
				fname := filepath.Join(output, string(rng.Name[:])+".bin")
				log.WithField("name", fname).Info("Creating")
				if err := os.WriteFile(fname, data, 0o644); err != nil {
					return fmt.Errorf("failed to write bundle data to file: %v", err)
				}
			}
		}
	} else {
		return fmt.Errorf("unsupported bundle type: %d", b.Type)
	}
	return nil
}

func Open(filename string) (*Bundle, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("failed to open file: %v", err)
	}
	bundle, err := Parse(f)
	if err != nil {
		f.Close()
		return nil, err
	}
	bundle.closer = f
	return bundle, nil
}

func (b *Bundle) Close() error {
	if b.closer != nil {
		err := b.closer.(io.Closer).Close()
		b.closer = nil
		return err
	}
	return nil
}

func Parse(r io.ReadSeeker) (*Bundle, error) {
	var bn Bundle
	bn.r = r

	if err := binary.Read(r, binary.LittleEndian, &bn.Header); err != nil {
		return nil, fmt.Errorf("failed to read bundle header: %v", err)
	}
	slices.Reverse(bn.Magic[:])

	if string(bn.Magic[:]) != Magic {
		return nil, fmt.Errorf("invalid magic: %s; expected 'BUND'", string(bn.Magic[:]))
	}

	switch bn.Type {
	case 3: // ExclaveCore
		var t3 Type3
		if err := binary.Read(r, binary.LittleEndian, &t3); err != nil {
			return nil, fmt.Errorf("failed to read bundle type 3: %v", err)
		}
		bn.TypeHeader = t3
		// parse footer/config
		if _, err := r.Seek(-int64(t3.FooterOffset), io.SeekEnd); err != nil {
			return nil, fmt.Errorf("failed to seek to bundle config data: %v", err)
		}

		fdata := make([]byte, t3.FooterSz)
		if _, err := r.Read(fdata); err != nil {
			return nil, fmt.Errorf("failed to read bundle data: %v", err)
		}

		if _, err := asn1.Unmarshal(fdata, &bn.Config); err != nil {
			return nil, fmt.Errorf("failed to ASN.1 parse bundle config: %v", err)
		}

		if err := bn.ParseFiles(); err != nil {
			return nil, fmt.Errorf("failed to parse bundle files: %v", err)
		}
	case 4: // AOP/DCP
		var t4 Type4
		if err := binary.Read(r, binary.LittleEndian, &t4); err != nil {
			return nil, fmt.Errorf("failed to read bundle type 4: %v", err)
		}
		bn.TypeHeader = t4
		for _, rng := range t4.Ranges {
			slices.Reverse(rng.Name[:])
			switch string(rng.Name[:]) {
			// case "uedt":
			// 	// parse device tree/config
			// 	log.WithField("name", string(rng.Name[:])).Debug("Device Tree")
			// 	if _, err := r.Seek(int64(rng.Offset), io.SeekStart); err != nil {
			// 		return nil, fmt.Errorf("failed to seek to bundle config data: %v", err)
			// 	}
			// 	dtdata := make([]byte, rng.Size)
			// 	if _, err := r.Read(dtdata); err != nil {
			// 		return nil, fmt.Errorf("failed to read bundle data: %v", err)
			// 	}
			// 	dt, err := devicetree.ParseData(bytes.NewReader(dtdata))
			// 	if err != nil {
			// 		os.WriteFile("uedt.bin", dtdata, 0o644)
			// 		return nil, fmt.Errorf("failed to parse device tree: %v", err)
			// 	}
			// 	log.Debug(dt.String())
			default:
				log.WithField("name", string(rng.Name[:])).Debug("unsupported")
			}
		}
	default:
		return nil, fmt.Errorf("unknown bundle type: %d", bn.Type)
	}

	return &bn, nil
}