package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"

	"github.com/elazarl/goproxy"

	"github.com/tidwall/sjson"
	"vaclive.party/software/lychee-slicer/patcher"
	"vaclive.party/software/lychee-slicer/platform"
)

const ip = "127.0.0.1"

type graphqlRequest struct {
	OperationName string `json:"operationName"`
}

type userData struct {
	operationName string
}

func main() {
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	program := flag.String("program", "", "program to patch and/or proxy")
	runPatch := flag.Bool("patch", false, "only patch the program")
	runProxy := flag.Bool("proxy", false, "only start the proxy and program")
	flag.Parse()
	if *program == "" {
		*program = platform.FindProgram()
		if *program == "" {
			flag.Usage()
			log.Fatal("Program not specified")
		}
	}

	if !*runPatch && !*runProxy {
		log.Println("No flags specified, running both patch and proxy")
		*runProxy = true
		*runPatch = true
	}

	if *runPatch {
		err := patcher.Patch(*program)
		if err != nil {
			log.Fatal("Patch:", err)
		}
	}

	if !*runProxy {
		return
	}

	log.Println("Starting proxy")
	proxy := goproxy.NewProxyHttpServer()
	proxy.Verbose = *verbose

	host := "graph.mango3d.io:443"
	proxy.OnRequest(goproxy.ReqHostIs(host)).HandleConnectFunc(goproxy.AlwaysMitm)
	proxy.OnRequest(goproxy.ReqHostIs(host)).DoFunc(func(req *http.Request, ctx *goproxy.ProxyCtx) (ret_req *http.Request, ret_res *http.Response) {
		defer func() {
			if e := recover(); e != nil {
				ctx.Warnf("Panic: %v", e)
				ret_req = req
				ret_res = goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusBadGateway, "Panic")
			}
		}()

		if req.Method != "POST" {
			return req, nil
		}

		body, err := io.ReadAll(req.Body)
		if err != nil {
			panic(err)
		}
		req.Body.Close()
		req.Body = io.NopCloser(bytes.NewBuffer(body))

		var graphql graphqlRequest
		if err := json.Unmarshal(body, &graphql); err != nil {
			ctx.Warnf("Unmarshal: %v", err)
			return req, nil
		}

		ctx.UserData = &userData{operationName: graphql.OperationName}

		fmt.Println("OperationName:", graphql.OperationName)
		return req, nil
	})

	modifyResponse := func(name string, callback func(string) string) {
		proxy.OnResponse(operationNameIs(name)).DoFunc(func(resp *http.Response, ctx *goproxy.ProxyCtx) (ret *http.Response) {
			defer func() {
				if e := recover(); e != nil {
					ret = goproxy.NewResponse(resp.Request, goproxy.ContentTypeText, http.StatusBadGateway, "Panic")
				}
			}()

			bb, err := io.ReadAll(resp.Body)
			if err != nil {
				panic(err)
			}
			resp.Body.Close()

			b := string(bb)
			log.Println(name+":", b)

			b = callback(b)

			resp.Body = io.NopCloser(bytes.NewBuffer([]byte(b)))
			return resp
		})
	}

	modifyResponse("GetUser", func(b string) string {
		b, _ = sjson.Set(b, "data.user.status.isPro", true)
		b, _ = sjson.Set(b, "data.user.status.isLifetime", true)
		b, _ = sjson.Set(b, "data.user.status.type", "lifetime")
		b, _ = sjson.Set(b, "data.user.status.featureTags", []string{"LYCHEE_SLICER_PRO_SLA", "LYCHEE_SLICER_PRO_FDM", "LYCHEE_SLICER_PREMIUM"})
		b, _ = sjson.Set(b, "data.user.licenses.0.macAddressesLimit", 9999999)
		return b
	})

	modifyResponse("userCheck", func(b string) string {
		b, _ = sjson.Set(b, "data.consumer.user.status.isPremium", true)
		b, _ = sjson.Set(b, "data.consumer.user.status.isPro", true)
		b, _ = sjson.Set(b, "data.consumer.user.status.isLifetime", true)
		b, _ = sjson.Set(b, "data.consumer.user.status.featureTags", []string{"LYCHEE_SLICER_PRO_SLA", "LYCHEE_SLICER_PRO_FDM", "LYCHEE_SLICER_PREMIUM"})
		b, _ = sjson.Set(b, "data.consumer.user.status.type", "lifetime")
		return b
	})

	modifyResponse("checkLicenseWithoutDuration", func(b string) string {
		b, _ = sjson.Set(b, "data.checkLicenseWithoutDuration.isPro", true)
		b, _ = sjson.Set(b, "data.checkLicenseWithoutDuration.isLifetime", true)
		b, _ = sjson.Set(b, "data.checkLicenseWithoutDuration.type", "lifetime")
		b, _ = sjson.Set(b, "data.checkLicenseWithoutDuration.mac_addresses_limit", 9999999)
		return b
	})

	modifyResponse("getLicense", func(b string) string {
		return b
	})

	modifyResponse("CanUserGetTrial", func(b string) string {
		return b
	})

	port, err := getFreePort()
	if err != nil {
		log.Fatal("GetFreePort:", err)
	}

	ln, err := net.Listen("tcp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		log.Fatal("Listen:", err)
	}

	log.Println("Listening on", ln.Addr().String(), port)

	go func() {
		cmd := exec.Command(*program, fmt.Sprintf("--proxy-server=%s:%d", ip, port), "--ignore-certificate-errors")
		cmd.Args = append(cmd.Args, flag.Args()...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		log.Fatal(cmd.Run())
		os.Exit(0)
	}()

	log.Fatal(http.Serve(ln, proxy))
}

func getFreePort() (port int, err error) {
	var a *net.TCPAddr
	if a, err = net.ResolveTCPAddr("tcp", "localhost:0"); err == nil {
		var l *net.TCPListener
		if l, err = net.ListenTCP("tcp", a); err == nil {
			defer l.Close()
			return l.Addr().(*net.TCPAddr).Port, nil
		}
	}
	return
}

func operationNameIs(operationName string) goproxy.RespCondition {
	return goproxy.RespConditionFunc(func(resp *http.Response, ctx *goproxy.ProxyCtx) bool {
		if ctx.UserData == nil {
			return false
		}
		userData := ctx.UserData.(*userData)
		return userData.operationName == operationName
	})
}