package signature

import (
	"encoding/json"
	"fmt"
	"maps"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/apex/log"
	"github.com/blacktop/go-macho"
	"github.com/blacktop/go-macho/types"
	"github.com/blacktop/ipsw/internal/utils"
	"github.com/blacktop/ipsw/pkg/disass"
	"github.com/blacktop/ipsw/pkg/kernelcache"
	"github.com/vbauerster/mpb/v8"
	"github.com/vbauerster/mpb/v8/decor"
)

func Parse(dir string) (sigs []Symbolicator, err error) {
	if err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() {
			if filepath.Ext(path) != ".json" {
				return nil
			}
			data, err := os.ReadFile(path)
			if err != nil {
				return err
			}
			var sig Symbolicator
			if err := json.Unmarshal(data, &sig); err != nil {
				return err
			}
			sigs = append(sigs, sig)
		}
		return nil
	}); err != nil {
		return nil, err
	}
	return sigs, nil
}

type SymbolMap map[uint64]string

func NewSymbolMap() SymbolMap {
	return make(SymbolMap)
}

func (sm SymbolMap) LoadJSON(infile string) error {
	data, err := os.ReadFile(infile)
	if err != nil {
		return fmt.Errorf("failed to read file: %v", err)
	}
	if err := json.Unmarshal(data, &sm); err != nil {
		return fmt.Errorf("failed to unmarshal json: %v", err)
	}
	return nil
}

func (sm SymbolMap) Add(addr uint64, symbol string) error {
	if sym, ok := sm[addr]; ok {
		if sym == symbol || sym == symbol+"_trap" {
			return nil // NOP
		}
		return fmt.Errorf("%#x already has symbol '%s', cannot add symbol '%s'", addr, sym, symbol)
	}
	sm[addr] = symbol
	return nil
}

func (sm SymbolMap) overwrite(addr uint64, symbol string) {
	sm[addr] = symbol
}

func (sm SymbolMap) Copy(m map[uint64]string) {
	maps.Copy(sm, m)
}

// Nearest returns the symbol whose address is the greatest key less than or
// equal to addr, along with that base address. ok is false when no symbol sits
// at or below addr. Callers can compute the offset as addr-base.
func (sm SymbolMap) Nearest(addr uint64) (base uint64, symbol string, ok bool) {
	for a, s := range sm {
		if a <= addr && (!ok || a > base) {
			base, symbol, ok = a, s, true
		}
	}
	return base, symbol, ok
}

func resolveTextSection(m *macho.File) (*types.Section, error) {
	for _, segment := range []string{"__TEXT_EXEC", "__TEXT"} {
		text := m.Section(segment, "__text")
		if text != nil {
			return text, nil
		}
	}
	return nil, fmt.Errorf("failed to find __text section in __TEXT_EXEC or __TEXT")
}

// symbolicateWithStats performs symbolication and returns (matched, missed, error).
// When quiet=true, verbose per-file logs are suppressed.
func (sm SymbolMap) symbolicateWithStats(m *macho.File, name string, sigs Symbolicator, quiet bool) (matched, missed int, err error) {
	seen := make(map[string]bool)

	text, err := resolveTextSection(m)
	if err != nil {
		return 0, 0, err
	}
	data, err := text.Data()
	if err != nil {
		return 0, 0, fmt.Errorf("failed to get data from resolved __text section: %v", err)
	}

	engine := disass.NewMachoDisass(m, &disass.Config{
		Data:         data,
		StartAddress: text.Addr,
		Quiet:        true,
	})

	if !quiet {
		log.WithField("name", name).Info("Analyzing MachO...")
	}
	if err := engine.Triage(); err != nil {
		return 0, 0, fmt.Errorf("first pass triage failed: %v", err)
	}

	cstrs, err := m.GetCStrings()
	if err != nil {
		return 0, 0, fmt.Errorf("failed to get cstrings: %v", err)
	}

	totalSigs := len(sigs.Signatures)
	matchedSigs := 0

	for _, sig := range sigs.Signatures {
		found := false
		for _, anchor := range sig.Anchors {
			if addr, ok := cstrs[fmt.Sprintf("%s.%s", anchor.Segment, anchor.Section)][anchor.String]; ok {
				if ok, loc := engine.Contains(addr); ok {
					fn, err := m.GetFunctionForVMAddr(loc)
					if err != nil {
						log.Errorf("failed to get function for address: %v", err)
						continue
					}
					if !quiet {
						utils.Indent(log.WithFields(log.Fields{
							"file":    name,
							"address": fmt.Sprintf("%#09x", fn.StartAddr),
							"symbol":  sig.Symbol,
						}).Info, 2)("Symbolicated")
					}
					if err := sm.Add(fn.StartAddr, sig.Symbol); err != nil {
						utils.Indent(log.WithError(err).Debug, 3)("failed to add to symbol map")
						// return fmt.Errorf("failed to add to symbol map: %v", err)
					}
					found = true
					matchedSigs++
					// attempt to symbolicate signature backtrace
					callerLoc := fn.StartAddr
					for _, caller := range sig.Backtrace {
						if !seen[caller] {
							if ok, loc := engine.Contains(callerLoc); ok {
								fcn, err := m.GetFunctionForVMAddr(loc)
								if err != nil {
									log.Errorf("failed to get function for address: %v", err)
									break // don't continue because we broke the caller chain (backtrace)
								}
								if !quiet {
									utils.Indent(log.WithFields(log.Fields{
										"file":    name,
										"address": fmt.Sprintf("%#09x", fn.StartAddr),
										"symbol":  caller,
									}).Info, 3)("Symbolicated (Caller)")
								}
								if err := sm.Add(fcn.StartAddr, caller); err != nil {
									utils.Indent(log.WithError(err).WithField("signature", sig.Symbol).Debug, 4)("failed to add 'caller' to symbol map")
									// return nil, fmt.Errorf("failed to add 'caller' to symbol map (for signature '%s'): %v", sig.Symbol, err)
								}
								callerLoc = fcn.StartAddr
							} else {
								if !quiet {
									utils.Indent(log.WithFields(log.Fields{
										"macho":  name,
										"caller": caller,
										"symbol": sig.Symbol,
									}).Warn, 2)("No XREFs to Caller found")
								}
								break
							}
							seen[caller] = true
						}
					}
					break // found symbol so break out of anchor loop
				} else {
					if !quiet {
						utils.Indent(log.WithFields(log.Fields{
							"macho":  name,
							"anchor": truncate(strconv.Quote(anchor.String), 40),
							"symbol": sig.Symbol,
						}).Warn, 2)("XREF Not Found For Anchor")
					}
				}
			} else {
				if !quiet {
					utils.Indent(log.WithFields(log.Fields{
						"macho":  name,
						"anchor": truncate(strconv.Quote(anchor.String), 40),
						"symbol": sig.Symbol,
					}).Debug, 3)("Anchor Not Found")
				}
			}
		}
		if !found {
			if !quiet {
				utils.Indent(log.WithFields(log.Fields{
					"macho":  name,
					"symbol": sig.Symbol,
				}).Warn, 2)("Signature Not Matched")
			}
		}
	}

	missedSigs := totalSigs - matchedSigs

	if !quiet {
		log.WithFields(log.Fields{
			"total":   totalSigs,
			"matched": matchedSigs,
			"missed":  missedSigs,
			"percent": fmt.Sprintf("%.4f%%", 100*float64(matchedSigs)/float64(totalSigs)),
		}).Info("📈 Symbolication STATS")
	}

	return matchedSigs, missedSigs, nil
}

func (sm SymbolMap) getSyscalls(m *macho.File) error {
	addr, syscalls, err := kernelcache.GetSyscallTable(m)
	if err != nil {
		return err
	}
	if err := sm.Add(addr, "bsd_syscall_table"); err != nil {
		utils.Indent(log.WithError(err).Debug, 2)("Adding bsd_syscall_table")
	}

	for _, syscall := range syscalls {
		if syscall.Name == "syscall" || syscall.Name == "enosys" {
			continue
		}
		if err := sm.Add(syscall.Call, syscall.Name); err != nil {
			utils.Indent(log.WithError(err).Debug, 2)("Adding syscall")
		}
	}

	return nil
}

func (sm SymbolMap) getMachTraps(m *macho.File) error {
	addr, machtraps, err := kernelcache.GetMachTrapTable(m)
	if err != nil {
		return fmt.Errorf("failed to get mach trap table: %v", err)
	}
	if err := sm.Add(addr, "mach_trap_table"); err != nil {
		utils.Indent(log.WithError(err).Debug, 2)("Adding mach_trap_table")
	}

	for _, machtrap := range machtraps {
		if machtrap.Name == "kern_invalid" {
			continue
		}
		if err := sm.Add(machtrap.Function, machtrap.Name+"_trap"); err != nil {
			utils.Indent(log.WithError(err).Debug, 2)("Adding mach_trap")
		}
	}

	return nil
}

func (sm SymbolMap) getMig(m *macho.File) error {
	addr, migs, err := kernelcache.GetMigSubsystems(m)
	if err != nil {
		return fmt.Errorf("failed to get MIG subsystems: %v", err)
	}
	if err := sm.Add(addr, "mig_kern_subsystem"); err != nil {
		utils.Indent(log.WithError(err).Debug, 2)("Adding mig_kern_subsystem")
	}

	for _, mig := range migs {
		if err := sm.Add(mig.KServer, strings.TrimSuffix(mig.Start.String(), "_subsystem")+"_server_routine"); err != nil {
			utils.Indent(log.WithError(err).Debug, 2)("Adding mig server_routine")
		}
		for idx, routine := range mig.Routines {
			if routine.KStubRoutine != 0 {
				if err := sm.Add(routine.KStubRoutine, mig.LookupRoutineName(idx)); err != nil {
					utils.Indent(log.WithError(err).Debug, 2)("Adding mig routine")
				}
			}
		}
	}

	return nil
}

func (sm SymbolMap) Symbolicate(infile string, sigs []Symbolicator, quiet bool) error {
	kc, err := macho.Open(infile)
	if err != nil {
		return fmt.Errorf("failed to open kernelcache: %v", err)
	}
	defer kc.Close()

	return sm.SymbolicateMachO(kc, filepath.Base(infile), sigs, quiet)
}

// SymbolicateMachO builds a kernel symbol map from built-in sources, optional
// signature packs, and recovered C++ metadata for an already-open Mach-O.
func (sm SymbolMap) SymbolicateMachO(kc *macho.File, kernelName string, sigs []Symbolicator, quiet bool) error {
	if kc == nil {
		return fmt.Errorf("nil kernelcache")
	}

	kv, err := kernelcache.GetVersion(kc)
	if err != nil {
		return fmt.Errorf("failed to get kernelcache version: %v", err)
	}

	if err := sm.getSyscalls(kc); err != nil {
		log.WithError(err).Warn("failed to get syscalls")
	}
	if err := sm.getMachTraps(kc); err != nil {
		log.WithError(err).Warn("failed to get mach traps")
	}
	if err := sm.getMig(kc); err != nil {
		log.WithError(err).Warn("failed to get MIG subsystems")
	}
	cppAdded, err := sm.addKernelCPPSymbols(kc, quiet)
	if err != nil {
		log.WithError(err).Warn("failed to discover C++ kernel symbols")
	}

	// Count valid signatures to set up progress bar
	var validSigs []Symbolicator
	for _, sig := range sigs {
		if ok, err := checkVersion(kv, sig); ok && err == nil {
			validSigs = append(validSigs, sig)
		}
	}

	// Set up progress bar for quiet mode
	var p *mpb.Progress
	var bar *mpb.Bar
	if quiet && len(validSigs) > 0 {
		// Use same blue color as apex/log info level for consistency
		blue := "\033[34m\033[1m"
		reset := "\033[0m"
		p = mpb.New(
			mpb.WithWidth(60),
			mpb.WithRefreshRate(180*time.Millisecond),
		)
		bar = p.AddBar(int64(len(validSigs)),
			mpb.BarFillerClearOnComplete(),
			mpb.PrependDecorators(
				decor.Name(blue+"   • Symbolicating signatures "+reset, decor.WC{C: decor.DindentRight}),
			),
			mpb.AppendDecorators(
				decor.CountersNoUnit("%d/%d", decor.WCSyncWidth),
			),
		)
	}

	processedSignature := false
	var totalMatched, totalMissed int

	for _, sig := range validSigs {
		// TODO: add support for OLD non-fileset KEXTs
		if kc.FileTOC.FileHeader.Type == types.MH_FILESET {
			m, err := kc.GetFileSetFileByName(sig.Target)
			if err != nil {
				if bar != nil {
					bar.Increment()
				}
				continue // fileset doesn't contain target
			}
			// symbolicate with signature
			matched, missed, err := sm.symbolicateWithStats(m, sig.Target, sig, quiet)
			if err != nil {
				if p != nil {
					p.Wait()
				}
				return err
			}
			totalMatched += matched
			totalMissed += missed
		} else {
			parts := strings.Split(sig.Target, ".")
			if kernelName != "" && len(parts) > 1 {
				// check if target macho file matches signature target
				if !strings.HasPrefix(strings.ToLower(kernelName), strings.ToLower(parts[len(parts)-1])) {
					if bar != nil {
						bar.Increment()
					}
					continue
				}
			}
			// symbolicate with signature
			matched, missed, err := sm.symbolicateWithStats(kc, sig.Target, sig, quiet)
			if err != nil {
				if p != nil {
					p.Wait()
				}
				return err
			}
			totalMatched += matched
			totalMissed += missed
		}

		if bar != nil {
			bar.Increment()
		}
		processedSignature = true
	}

	if p != nil {
		p.Wait()
	}

	// Print summary stats in quiet mode
	if quiet && processedSignature {
		total := totalMatched + totalMissed
		if total > 0 {
			log.WithFields(log.Fields{
				"matched": totalMatched,
				"missed":  totalMissed,
				"percent": fmt.Sprintf("%.2f%%", 100*float64(totalMatched)/float64(total)),
			}).Info("📈 Symbolication complete")
		}
	}

	switch {
	case len(validSigs) == 0 && len(sigs) > 0:
		log.Warn("No valid signatures matched kernelcache version; using built-in and C++ symbols only")
	case len(validSigs) > 0 && !processedSignature:
		log.Warn("No signatures matched in kernelcache; using built-in and C++ symbols only")
	case len(sigs) == 0 && cppAdded == 0 && !quiet:
		log.Warn("No signature pack supplied; emitted only built-in kernel symbols")
	}

	return nil
}
