package wrapper

import (
	"fmt"
	"io/ioutil"
	"os"

	"github.com/alecthomas/kong"
	log "github.com/sirupsen/logrus"
	"google.golang.org/protobuf/proto"

	"ios-wrapper/internal/wrapper/addr2line"
	. "ios-wrapper/internal/wrapper/ofile"
	"ios-wrapper/pkg/ios/fat"
	"ios-wrapper/pkg/protomodel"
)

const envLogLevel = "LOG_LEVEL"
const defaultLogLevel = log.InfoLevel

func getLogLevel() log.Level {
	levelString, exists := os.LookupEnv(envLogLevel)
	if !exists {
		return defaultLogLevel
	}

	level, err := log.ParseLevel(levelString)
	if err != nil {
		return defaultLogLevel
	}

	return level
}

func Cli() {
	var cli Argument
	ctx := kong.Parse(&cli)
	command := ctx.Selected().Name

	log.SetLevel(getLogLevel())

	if command == "info" {
		arg := cli.Info
		printOFile(arg.OFile)
		return
	} else if command == "signed-bcell" {
		arg := cli.SignedBcell
		makeSignedBcell(arg.Out, arg.Bcell)
		return
	} else if command == "display-bcell" {
		arg := cli.DisplayBcell
		displayBcell(arg.Bcell)
		return
	} else if command == "addr2line" {
		arg := cli.Addr2Line
		resolveAddresses(arg.Dwarf, arg.Load, arg.Addresses)
		return
	} else if command == "lipo split" {
		arg := cli.Lipo.Split
    fat.FatSplit(arg.Fat)
    return
  } else if command == "lipo join" {
		arg := cli.Lipo.Join
    fat.FatJoin(arg.Macho, arg.Out)
    return
  }

	var pc ProgramContext
	var ofile OFile = nil

	switch command {
	case "wrap":
		arg := cli.Wrap
		ofile = NewOFile(arg.OFile)
		pc.remove_codesign = true
		pc.strip_init_pointers = true
		pc.dylib_to_add = []string{"@rpath/bcell.framework/bcell"}
		pc.rpath_to_add = []string{"@executable_path/Frameworks"}
		pc.outfile = arg.Out
		pc.bcellfile = arg.Bcell
		pc.signed = arg.Signed
		pc.ReadUserConfig(arg.Config)

  case "vltk":
		arg := cli.Vltk
		ofile = NewOFile(arg.OFile)
		pc.remove_codesign = true
		pc.dylib_to_add = []string{"@rpath/GTJetModule.framework/GTJetModule"}
		pc.outfile = arg.Out

	case "remove-codesign", "remove-signature":
		arg := cli.RemoveCodesign
		ofile = NewOFile(arg.OFile)
		pc.remove_codesign = true
		pc.outfile = arg.Out

	default:
		return
	}

	pc.Process(ofile)
	ofile.Cleanup()
}

func printOFile(ofile string) {
	f, _ := os.OpenFile(ofile, os.O_RDWR, 0644)
	printer := InfoPrinterFromFile(f)
	printer.Print()
}

func makeSignedBcell(outfile string, infile string) {
	raw_data, err := ioutil.ReadFile(infile)
	if err != nil {
		log.Panic("Cannot read bcell.dat")
	}
	data := &protomodel.BcellFile{}
	err = proto.Unmarshal(raw_data, data)
	if err != nil {
		log.Panic("Invalid bcell.dat")
	}

	blocks := []*protomodel.Block{
		{
			Key:   0,
			Value: raw_data,
		},
	}
	signed_data := protomodel.SignedData{
		Blocks: blocks,
	}
	signed_data_b, err := proto.Marshal(&signed_data)
	err = ioutil.WriteFile(outfile, signed_data_b, 0644)
	if err != nil {
		log.Panic("Cannot write SignedData to file")
	}
}

func displayBcell(bfile string) {
	raw_data, err := ioutil.ReadFile(bfile)
	if err != nil {
		log.Panic("Invalid Protobuf<BcellFile> bcell.dat (1)")
	}
	data := &protomodel.BcellFile{}
	err = proto.Unmarshal(raw_data, data)
	if err != nil {
		log.Panic("Invalid Protobuf<BcellFile> bcell.dat (2)")
	}

	fmt.Printf("[+] User Config: %+v\n", data.BcellConfig)
	for arch, info := range data.MachoInfos {
		fmt.Printf("[+] Arch %s:\n", arch)
		fmt.Printf(" | PointerSize  : %+v\n", info.PointerSize)
		fmt.Printf(" | Image Base   : 0x%x\n", info.ImageBase)
		fmt.Printf(" | Init Pointers:\n")
		for _, init_ptr := range info.InitPointers {
			fmt.Printf(
				" |   offset 0x%x => addr 0x%x\n",
				init_ptr.Offset,
				init_ptr.Value,
			)
		}
	}
}

func resolveAddresses(dwarf string, load string, addresses []string) {
	mf, valid_macho := NewOFile(dwarf).(*MachoFile)
	if !valid_macho {
		log.Error("Input DWARF file is not a valid Macho binary")
		return
	}

	loadaddr, err := addr2line.ParseAddressString(load)
	if err != nil {
		log.Error("Load address is not a valid hex number (%s)", load)
		return
	}

	resolver := addr2line.NewResolver(mf.Context(), loadaddr, addresses)
	for ; ; resolver.HasNext() {
		resolved, err := resolver.Next()
		if err != nil {
			fmt.Printf("[?] Error: %s\n", err)
		} else if resolved.Valid() {
			fmt.Printf("[+] Found 0x%x => 0x%x %s %s:%d\n",
				resolved.Raw, resolved.Base, resolved.Symbol, resolved.File, resolved.Line)
		} else {
			fmt.Printf("[-] Not found 0x%x\n", resolved.Raw)
		}
	}
}