package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/go-rod/rod"
	"github.com/go-rod/rod/lib/launcher"
	"github.com/go-rod/rod/lib/proto"
	"github.com/go-rod/stealth"
)

var (
	globalBrowser *rod.Browser

	statsMu       sync.Mutex
	totalRequests int64
	totalErrors   int64
	activeTabs    int64
	startTime     time.Time

	tmpDir string
)

type AdvancedAction struct {
	Type     string `json:"type"`
	Selector string `json:"selector,omitempty"`
	Value    string `json:"value,omitempty"`
	FileName string `json:"fileName,omitempty"`
}

type Payload struct {
	Timeout           int                    `json:"timeout,omitempty"`
	Width             int                    `json:"width,omitempty"`
	Height            int                    `json:"height,omitempty"`
	UserAgent         string                 `json:"userAgent,omitempty"`
	IgnoreHTTPSErrors bool                   `json:"ignoreHTTPSErrors,omitempty"`
	DisableJS         bool                   `json:"disableJS,omitempty"`
	BlockImages       bool                   `json:"blockImages,omitempty"`
	SniffNetwork      bool                   `json:"sniffNetwork,omitempty"`
	MaxSniffSize      int                    `json:"maxSniffSize,omitempty"`
	Headers           map[string]interface{} `json:"headers,omitempty"`
	Actions           []AdvancedAction       `json:"actions"`
}

type StandardResponse struct {
	Logs   []string               `json:"logs"`
	Result map[string]interface{} `json:"result"`
	Error  string                 `json:"error,omitempty"`
}

type StatsResponse struct {
	Uptime           string `json:"uptime"`
	TotalRequests    int64  `json:"totalRequests"`
	TotalErrors      int64  `json:"totalErrors"`
	ActiveTabs       int64  `json:"activeTabs"`
	OpenPagesCount   int    `json:"openPagesCount"`
	GoVersion        string `json:"goVersion"`
	NumGoroutine     int    `json:"numGoroutine"`
	AllocMemoryMB    uint64 `json:"allocMemoryMB"`
	TotalAllocMB     uint64 `json:"totalAllocMB"`
	SysMemoryMB      uint64 `json:"sysMemoryMB"`
	NumGC            uint32 `json:"numGC"`
}

type SniffedNetworkActivity struct {
	URL             string            `json:"url"`
	Method          string            `json:"method"`
	RequestHeaders  map[string]string `json:"requestHeaders,omitempty"`
	ResponseStatus  int               `json:"responseStatus,omitempty"`
	ResponseHeaders map[string]string `json:"responseHeaders,omitempty"`
	ResponseBody    string            `json:"responseBody,omitempty"`
}

func main() {
	port := "7860"
	startTime = time.Now()

	tmpDir = os.TempDir()

	initGlobalBrowser()
	defer globalBrowser.MustClose()

	http.HandleFunc("/", statsHandler)
	http.HandleFunc("/execute", executeHandler)
	// http.HandleFunc("/stats", statsHandler)

	fs := http.FileServer(http.Dir(tmpDir))
	http.Handle("/tmp/", http.StripPrefix("/tmp/", fs))

	fmt.Printf("Running on port %s...\n", port)
	if err := http.ListenAndServe(":"+port, nil); err != nil {
		log.Fatal(err)
	}
}

func initGlobalBrowser() {
	path, _ := launcher.LookPath()
	u := launcher.New().
		Bin(path).
		NoSandbox(true).
		Headless(true).
		Devtools(false).
		Set("disable-gpu", "true").
		Set("disable-dev-shm-usage", "true").
		Set("disable-blink-features", "AutomationControlled").
		Set("exclude-switches", "enable-automation").
		Set("disable-gl-drawing-for-tests", "true").
		Set("disable-software-rasterizer", "true").
		Set("disable-canvas-aa", "true").
		Set("disable-2d-canvas-clip-antialiasing", "true").
		Set("disable-gpu-rasterization", "true").
		Set("no-zygote", "true").
		Set("single-process", "true").
		Set("disable-infobars", "true").
		Set("disable-notifications", "true").
		Set("disable-popup-blocking", "true").
		Set("disable-background-networking", "true").
		Set("disable-renderer-backgrounding", "true").
		Set("no-first-run", "true").
		Set("no-default-browser-check", "true").
		Set("window-size", "1366,768").
		Set("lang", "en-US,en").
		Set("platform", "Win32").
		Set("blink-settings", "primaryHoverType=2,availableHoverTypes=2,primaryPointerType=4,availablePointerTypes=4").
		MustLaunch()

	globalBrowser = rod.New().ControlURL(u).MustConnect()
}

func executeHandler(w http.ResponseWriter, r *http.Request) {
	statsMu.Lock()
	totalRequests++
	activeTabs++
	statsMu.Unlock()

	defer func() {
		statsMu.Lock()
		activeTabs--
		statsMu.Unlock()
	}()

	if r.Method != http.MethodPost {
		http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
		return
	}

	var req Payload
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	tabTimeout := 30
	if req.Timeout > 0 {
		tabTimeout = req.Timeout
	}

	var logs []string
	results := make(map[string]interface{})
	
	var screenshots []string
	var evals []interface{}

	var networkMu sync.Mutex
	var finalNetworkLogs []SniffedNetworkActivity

	pageCtx, cancelPage := context.WithTimeout(context.Background(), time.Duration(tabTimeout)*time.Second)
	defer cancelPage()

	scopedBrowser := globalBrowser.Context(pageCtx)
	p := stealth.MustPage(scopedBrowser)
	defer p.MustClose()

	_ = p.MustSetExtraHeaders("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8")
	_ = p.MustSetExtraHeaders("Accept-Language", "en-US,en;q=0.9")
	_ = p.MustSetExtraHeaders("Sec-Ch-Ua", `"Chromium";v="124", "Not A(Brand";v="24", "Google Chrome";v="124"`)
	_ = p.MustSetExtraHeaders("Sec-Ch-Ua-Mobile", "?0")
	_ = p.MustSetExtraHeaders("Sec-Ch-Ua-Platform", `"Windows"`)

	width := 1366
	height := 768
	if req.Width > 0 && req.Height > 0 {
		width = req.Width
		height = req.Height
	}
	_ = p.SetViewport(&proto.EmulationSetDeviceMetricsOverride{
		Width:  width,
		Height: height,
	})

	uaStr := "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36"
	if req.UserAgent != "" {
		uaStr = req.UserAgent
	}
	_ = p.SetUserAgent(&proto.NetworkSetUserAgentOverride{
		UserAgent: uaStr,
	})

	if req.SniffNetwork {
		maxSize := 1024 * 50
		if req.MaxSniffSize > 0 {
			maxSize = req.MaxSniffSize
		}

		router := p.HijackRequests()
		router.MustAdd("*", func(ctx *rod.Hijack) {
			ctx.MustLoadResponse()
			
			res := ctx.Response
			reqInfo := ctx.Request
			
			contentType := strings.ToLower(res.Headers().Get("Content-Type"))
			isData := strings.Contains(contentType, "json") || 
			          strings.Contains(contentType, "text/html") || 
			          strings.Contains(contentType, "javascript")

			activity := SniffedNetworkActivity{
				URL:            reqInfo.URL().String(),
				Method:         reqInfo.Method(),
				ResponseStatus: 200,
			}

			reqHeadersMentah, _ := json.Marshal(reqInfo.Headers())
			resHeadersMentah, _ := json.Marshal(res.Headers())

			reqHeadersMap := make(map[string]string)
			resHeadersMap := make(map[string]string)

			_ = json.Unmarshal(reqHeadersMentah, &reqHeadersMap)
			_ = json.Unmarshal(resHeadersMentah, &resHeadersMap)

			activity.RequestHeaders = reqHeadersMap
			activity.ResponseHeaders = resHeadersMap

			if isData {
				body := ctx.Response.Body()
				if len(body) > maxSize {
					activity.ResponseBody = body[:maxSize] + " ... [TRUNCATED]"
				} else {
					activity.ResponseBody = body
				}
			} else {
				activity.ResponseBody = "[SKIPPED: NON-DATA OR LARGE ASSET]"
			}

			networkMu.Lock()
			finalNetworkLogs = append(finalNetworkLogs, activity)
			networkMu.Unlock()
			
			ctx.ContinueRequest(&proto.FetchContinueRequest{})
		})
		go router.Run()
	}

	_, _ = p.Eval(`
		(() => {
			const tamperWebdriver = () => {
				try {
					Object.defineProperty(navigator, 'webdriver', {
						get: () => false,
						configurable: false,
						enumerable: true
					});
				} catch (e) {}
			};

			tamperWebdriver();

			if (!window.chrome) {
				Object.defineProperty(window, 'chrome', {
					configurable: true,
					enumerable: true,
					value: { runtime: {} }
				});
			}

			window.__turnstileToken = '';
			const hookTurnstile = () => {
				if (!window.turnstile || window.turnstile.__hooked) return;
				const origRender = window.turnstile.render.bind(window.turnstile);
				window.turnstile.render = function(container, params) {
					if (params && typeof params.callback === 'function') {
						const origCb = params.callback;
						params.callback = function(token) {
							window.__turnstileToken = token;
							return origCb(token);
						};
					}
					return origRender(container, params);
				};
				window.turnstile.__hooked = true;
			};
			Object.defineProperty(window, 'turnstile', {
				configurable: true,
				set(v) { this.__ts = v; setTimeout(hookTurnstile, 0); },
				get()  { return this.__ts; }
			});
		})();
	`)

	for i, act := range req.Actions {
		logMsg := fmt.Sprintf("Step %d: [%s]", i+1, act.Type)
		if act.Selector != "" {
			logMsg += fmt.Sprintf(" -> Target: %s", act.Selector)
		}
		logs = append(logs, logMsg)

		var err error

		switch act.Type {
		case "navigate":
			err = p.Navigate(act.Value)
			if err == nil {
				_ = p.WaitDOMStable(500*time.Millisecond, 0.1)
				if req.BlockImages {
					injectStyleScript := `const style = document.createElement('style'); style.innerText = 'img, svg, picture, [style*="background-image"] { display: none !important; }'; document.head.appendChild(style);`
					_, _ = p.Eval(injectStyleScript)
				}
			}

		case "wait":
			_, err = p.Element(act.Selector)

		case "click":
			var el *rod.Element
			el, err = p.Element(act.Selector)
			if err == nil {
				err = el.Click(proto.InputMouseButtonLeft, 1)
			}

		case "type":
			var el *rod.Element
			el, err = p.Element(act.Selector)
			if err == nil {
				err = el.Input(act.Value)
			}

		case "clear":
			var el *rod.Element
			el, err = p.Element(act.Selector)
			if err == nil {
				err = el.SelectAllText()
				if err == nil {
					err = el.Input("")
				}
			}

		case "hover":
			var el *rod.Element
			el, err = p.Element(act.Selector)
			if err == nil {
				err = el.Hover()
			}

		case "select":
			var el *rod.Element
			el, err = p.Element(act.Selector)
			if err == nil {
				err = el.Select([]string{act.Value}, true, rod.SelectorTypeRegex)
			}

		case "key":
			jsInputScript := fmt.Sprintf(`(() => { const activeEl = document.activeElement; if (activeEl) { activeEl.value += %q; activeEl.dispatchEvent(new Event('input', { bubbles: true })); activeEl.dispatchEvent(new Event('change', { bubbles: true })); } })()`, act.Value)
			_, err = p.Eval(jsInputScript)

		case "scroll":
			if act.Selector != "" {
				var el *rod.Element
				el, err = p.Element(act.Selector)
				if err == nil {
					err = el.ScrollIntoView()
				}
			} else {
				valInt, _ := strconv.Atoi(act.Value)
				_, err = p.Eval(fmt.Sprintf("window.scrollTo(0, %d);", valInt))
			}

		case "setLocation":
			coords := strings.Split(act.Value, ",")
			lat, _ := strconv.ParseFloat(coords[0], 64)
			long, _ := strconv.ParseFloat(coords[1], 64)
			accuracy := float64(100)
			_ = p.SetViewport(&proto.EmulationSetDeviceMetricsOverride{
				Width:  width,
				Height: height,
			})
			err = proto.EmulationSetGeolocationOverride{
				Latitude:  &lat,
				Longitude: &long,
				Accuracy:  &accuracy,
			}.Call(p)

		case "evaluate":
			var jsResult *proto.RuntimeRemoteObject
			jsResult, err = p.Eval(act.Value)
			if err == nil && jsResult != nil {
				evals = append(evals, jsResult.Value.Val())
			}

		case "extractTurnstile":
			var iframe *rod.Element
			for retry := 0; retry < 15; retry++ {
				iframe, _ = p.Element(`iframe[src*="challenges.cloudflare.com"], iframe[title*="verification"], .cf-turnstile`)
				if iframe != nil {
					break
				}
				time.Sleep(500 * time.Millisecond)
			}

			if iframe != nil {
				_ = iframe.ScrollIntoView()
				err = iframe.Click(proto.InputMouseButtonLeft, 1)
			}

			var tokenVal *proto.RuntimeRemoteObject
			harvestScript := `
				(() => {
					if (window.__turnstileToken && window.__turnstileToken.length > 10) return window.__turnstileToken;
					let input = document.querySelector('input[name="cf-turnstile-response"]') || 
					            document.querySelector('input[name="g-recaptcha-response"]') ||
					            document.getElementById('cf-turnstile-response');
					return input ? input.value : "";
				})()
			`
			for retry := 0; retry < 30; retry++ {
				tokenVal, err = p.Eval(harvestScript)
				if err == nil && tokenVal != nil && tokenVal.Value.String() != "" {
					break
				}
				time.Sleep(700 * time.Millisecond)
			}

			if tokenVal != nil && tokenVal.Value.String() != "" {
				results["turnstileToken"] = tokenVal.Value.String()
			} else {
				results["turnstileToken"] = "FAILED_OR_TIMEOUT"
			}

		case "screenshot":
			var buf []byte
			if act.Selector != "" {
				var el *rod.Element
				el, err = p.Element(act.Selector)
				if err == nil {
					buf, err = el.Screenshot(proto.PageCaptureScreenshotFormatPng, 100)
				}
			} else {
				buf = p.MustScreenshot("")
			}

			if err == nil && len(buf) > 0 {
				name := act.FileName
				if name == "" {
					name = fmt.Sprintf("ss-%d.png", time.Now().UnixNano())
				}
				name = strings.ReplaceAll(name, " ", "_")
				filePath := filepath.Join(tmpDir, name)
				err = os.WriteFile(filePath, buf, 0644)
				if err == nil {
					screenshots = append(screenshots, "/tmp/"+name)
				}
			}

		case "screenshotFullPage":
			buf := p.MustScreenshotFullPage()
			if len(buf) > 0 {
				name := act.FileName
				if name == "" {
					name = fmt.Sprintf("fullss-%d.png", time.Now().UnixNano())
				}
				name = strings.ReplaceAll(name, " ", "_")
				filePath := filepath.Join(tmpDir, name)
				err = os.WriteFile(filePath, buf, 0644)
				if err == nil {
					screenshots = append(screenshots, "/tmp/"+name)
				}
			}
		}

		if err != nil {
			statsMu.Lock()
			totalErrors++
			statsMu.Unlock()
			returnError(w, logs, fmt.Errorf("failed in step %d (%s): %w", i+1, act.Type, err))
			return
		}
	}

	results["screenshots"] = screenshots
	results["evals"] = evals
	if req.SniffNetwork {
		results["networkLogs"] = finalNetworkLogs
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(StandardResponse{
		Logs:   logs,
		Result: results,
	})
}

func statsHandler(w http.ResponseWriter, r *http.Request) {
	var m runtime.MemStats
	runtime.ReadMemStats(&m)

	pagesCount := 0
	if globalBrowser != nil {
		pages, err := globalBrowser.Pages()
		if err == nil {
			pagesCount = len(pages)
		}
	}

	statsMu.Lock()
	res := StatsResponse{
		Uptime:         time.Since(startTime).String(),
		TotalRequests:  totalRequests,
		TotalErrors:    totalErrors,
		ActiveTabs:     activeTabs,
		OpenPagesCount: pagesCount,
		GoVersion:      runtime.Version(),
		NumGoroutine:   runtime.NumGoroutine(),
		AllocMemoryMB:  m.Alloc / 1024 / 1024,
		TotalAllocMB:   m.TotalAlloc / 1024 / 1024,
		SysMemoryMB:    m.Sys / 1024 / 1024,
		NumGC:          m.NumGC,
	}
	statsMu.Unlock()

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(res)
}

func returnError(w http.ResponseWriter, logs []string, err error) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusInternalServerError)
	json.NewEncoder(w).Encode(StandardResponse{
		Logs:  logs,
		Error: err.Error(),
	})
}