package wrapper

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

	"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
	} else if command == "bcell2header" {
		arg := cli.BcellToHeader
		bcell2header(arg.Bcell, 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.remove_inits = 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

	case "remove-imports":
		arg := cli.RemoveImports
		ofile = NewOFile(arg.OFile)
		pc.remove_imports = true
		pc.outfile = arg.Out
		pc.bcellfile = arg.Bcell

	case "pepe":
		arg := cli.Pepe
		ofile = NewOFile(arg.OFile)
		if arg.FullRemoval {
			pc.remove_exports = true
			pc.remove_symbol_table = true
			pc.remove_imports = true
			pc.remove_inits = true
			pc.remove_codesign = true
			pc.remove_others = true
		}
		pc.remove_imports = arg.RemoveBindSymbols
		pc.remove_codesign = arg.RemoveCodeSign
		pc.remove_inits = arg.RemoveInitFunctions
		pc.remove_others = arg.RemoveOthers
		pc.remove_exports = arg.RemoveExports
		pc.remove_symbol_table = arg.RemoveSymbolTable
		pc.dylib_to_add = arg.Dylibs
		pc.rpath_to_add = arg.Rpath
		pc.outfile = arg.Out
		pc.bcellfile = arg.Bcell
		pc.symbols_keep = arg.KeepImports

	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,
			)
		}
		fmt.Printf(" | Bind Symbols:\n")
		// for _, symbol := range info.Symbols {
		// 	lib := strings.Replace(symbol.Libname, "/System/Library/Frameworks/", "", 1)
		// 	fmt.Printf(" |   %s offset=0x%x segmentID=0x%x\n", symbol.Name, symbol.Offset, symbol.Segment)
		// 	fmt.Printf(" |      from=%s\n", lib)
		// }
	}
}

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

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

	f, err := os.Create(header)
	if err != nil {
		log.Panic("Cannot open header file for writing")
	}
	defer f.Close()

	w := bufio.NewWriter(f)

	// fmt.Printf("[+] User Config: %+v\n", data.BcellConfig)
	fmt.Fprintf(w, "#include<stdint.h>\n")
	fmt.Fprintf(w, "namespace bshield_data{\n")
	for arch, info := range data.MachoInfos {
		fmt.Fprintf(w, "const char* arch = \"%s\";\n", arch)
		fmt.Fprintf(w, "unsigned int pointer_size = %d;\n", info.PointerSize)
		fmt.Fprintf(w, "uint64_t image_base = 0x%x;\n", info.ImageBase)
		fmt.Fprintf(w, "uint64_t main = 0x%x;\n", info.Main)

		fmt.Fprintf(w, "struct init_pointer {uint64_t offset; uint64_t value;};\n")
		fmt.Fprintf(w, "int num_init_pointers = %d;\n", len(info.InitPointers))
		fmt.Fprintf(w, "struct init_pointer init_pointers_offsets[] = {\n")
		for _, init_ptr := range info.InitPointers {
			fmt.Fprintf(w, "  {0x%x, 0x%x},\n", init_ptr.Offset, init_ptr.Value)
		}
		fmt.Fprintf(w, "};\n")

		fmt.Fprintf(w, "__attribute__((section(\"__DATA,bshield\")))\n")
		fmt.Fprintf(w, "char libs[] =\n")
		for _, lib := range info.Symbols.Libs {
			fmt.Fprintf(w, "  \"%s\\0\"\n", lib)
		}
		fmt.Fprintf(w, ";\n")

		fmt.Fprintf(w, "__attribute__((section(\"__DATA,bshield\")))\n")
		fmt.Fprintf(w, "char symbols[] =\n")
		for _, symbol := range info.Symbols.Symbols {
			fmt.Fprintf(w, "  \"%s\\0\"\n", symbol)
		}
		fmt.Fprintf(w, ";\n")

		fmt.Fprintf(w, "// very compact symbol table,\n")
		fmt.Fprintf(w, "// [lib idx/*4 bytes*/, nsymbol/*4 byte*/]\n")
		fmt.Fprintf(w, "// repeat nsymbol times [name offset/*3 bytes*/, segment idx/**/, offset /*4 btyes*/]\n")
		fmt.Fprintf(w, "// name offset is 3 bytes because we don't think we should have a table size > 2^(3 * 8)\n")

		fmt.Fprintf(w, "__attribute__((section(\"__DATA,bshield\")))\n")
		fmt.Fprintf(w, "uint32_t encoded_table[] = {\n")
		n_instructions := 0
		for i, table := range info.Symbols.Tables {
			fmt.Fprintf(w, "  // %s\n", info.Symbols.Libs[i])
			fmt.Fprintf(w, "  %d/*lib offset*/,\n", table.LibIndex)
			fmt.Fprintf(w, "  %d/*nsymbols*/,\n", table.Nsymbols)
			n_instructions += 2
			for _, symbol := range table.Symbols {
				fmt.Fprintf(w, "  %d, 0x%x,\n", (symbol.SymbolIndex<<8)|symbol.SegmentIndex, symbol.Offset)
				n_instructions += 2
			}
			fmt.Fprintf(w, "\n")
		}
		fmt.Fprintf(w, "};\n")
		fmt.Fprintf(w, "uint32_t n_instructions = %d;\n", n_instructions)

		fmt.Fprintf(w, "__attribute__((section(\"__DATA,bshield\")))\n")
		fmt.Fprintf(w, "uint32_t special_selectors_idx[] = {\n")
		for _, selector := range info.GetSpecialSelectors() {
			fmt.Fprintf(w, "%x,\n", selector.Idx)
		}
		fmt.Fprintf(w, "};\n")

		fmt.Fprintf(w, "__attribute__((section(\"__DATA,bshield\")))\n")
		fmt.Fprintf(w, "const char* special_selectors_name[] = {\n")
		for _, selector := range info.GetSpecialSelectors() {
			fmt.Fprintf(w, "\"%s\",\n", selector.Name)
		}
		fmt.Fprintf(w, "};\n")
		fmt.Fprintf(w, "uint32_t n_selectors = %d;\n", len(info.GetSpecialSelectors()))
	}
	fmt.Fprintf(w, "}// namespace bshield_data\n")
	w.Flush()
}