package blocker

import (
	"log"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"
	"unsafe"
)

var (
	user32                       = syscall.NewLazyDLL("user32.dll")
	procEnumWindows              = user32.NewProc("EnumWindows")
	procIsWindowVisible          = user32.NewProc("IsWindowVisible")
	procGetWindowTextW           = user32.NewProc("GetWindowTextW")
	procGetWindowTextLengthW     = user32.NewProc("GetWindowTextLengthW")
	procGetWindowThreadProcessId = user32.NewProc("GetWindowThreadProcessId")
	procGetWindow                = user32.NewProc("GetWindow")
	procGetWindowLongW           = user32.NewProc("GetWindowLongW")

	dwmapi                    = syscall.NewLazyDLL("dwmapi.dll")
	procDwmGetWindowAttribute = dwmapi.NewProc("DwmGetWindowAttribute")

	kernel32 = syscall.NewLazyDLL("kernel32.dll")
)

const (
	GW_OWNER         = 4
	WS_EX_TOOLWINDOW = 0x00000080
	DWMWA_CLOAKED    = 14
)

func isRealGUIWindow(hwnd uintptr) bool {
	// 1. Phải đang hiển thị
	ret, _, _ := procIsWindowVisible.Call(hwnd)
	if ret == 0 {
		return false
	}

	// 2. Không được có chủ sở hữu (phải là cửa sổ chính - top-level)
	owner, _, _ := procGetWindow.Call(hwnd, GW_OWNER)
	if owner != 0 {
		return false
	}

	// 3. Không phải là tool window (WS_EX_TOOLWINDOW)
	gwlExStyle := int32(-20) // GWL_EXSTYLE = -20
	style, _, _ := procGetWindowLongW.Call(hwnd, uintptr(gwlExStyle))
	if (style & WS_EX_TOOLWINDOW) != 0 {
		return false
	}

	// 4. Không bị cloaked bởi DWM (ví dụ: app UWP bị treo/chạy ngầm, màn hình ảo)
	var cloaked uint32
	hr, _, _ := procDwmGetWindowAttribute.Call(hwnd, DWMWA_CLOAKED, uintptr(unsafe.Pointer(&cloaked)), 4)
	if hr == 0 && cloaked != 0 {
		return false
	}

	return true
}

type WindowInfo struct {
	PID         uint32
	Title       string
	ProcessName string
}

type Blocker struct {
	mu              sync.Mutex
	allowedKeywords []string
	running         bool
	stopChan        chan struct{}
	OnBlocked       func(procName string, title string)
	OnKill          func(procName string, title string)
}

var Instance = &Blocker{
	allowedKeywords: []string{"chrome", "idea64", "vscode", "wails", "simple_care", "client"},
	stopChan:        make(chan struct{}),
}

func getProcessMap() (map[uint32]string, error) {
	snapshot, err := syscall.CreateToolhelp32Snapshot(syscall.TH32CS_SNAPPROCESS, 0)
	if err != nil {
		return nil, err
	}
	defer syscall.CloseHandle(snapshot)

	var pe syscall.ProcessEntry32
	pe.Size = uint32(unsafe.Sizeof(pe))

	err = syscall.Process32First(snapshot, &pe)
	if err != nil {
		return nil, err
	}

	pm := make(map[uint32]string)
	for {
		name := syscall.UTF16ToString(pe.ExeFile[:])
		pm[pe.ProcessID] = name

		err = syscall.Process32Next(snapshot, &pe)
		if err != nil {
			break
		}
	}
	return pm, nil
}

func getWindowText(hwnd uintptr) string {
	length, _, _ := procGetWindowTextLengthW.Call(hwnd)
	if length == 0 {
		return ""
	}
	buf := make([]uint16, length+1)
	procGetWindowTextW.Call(hwnd, uintptr(unsafe.Pointer(&buf[0])), length+1)
	return syscall.UTF16ToString(buf)
}

var (
	enumWindowsMutex sync.Mutex
	enumWindowsList  []WindowInfo
	enumProcessMap   map[uint32]string
)

var enumWindowsCallback = syscall.NewCallback(func(hwnd uintptr, lParam uintptr) uintptr {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("[BLOCKER] Callback panic recovered: %v", r)
		}
	}()

	if !isRealGUIWindow(hwnd) {
		return 1
	}

	title := getWindowText(hwnd)
	if title == "" {
		return 1
	}

	var pid uint32
	procGetWindowThreadProcessId.Call(hwnd, uintptr(unsafe.Pointer(&pid)))

	procName := enumProcessMap[pid]
	if procName == "" {
		procName = "Unknown"
	}

	enumWindowsList = append(enumWindowsList, WindowInfo{
		PID:         pid,
		Title:       title,
		ProcessName: procName,
	})

	return 1
})

func EnumerateGUIWindows() ([]WindowInfo, error) {
	pMap, err := getProcessMap()
	if err != nil {
		return nil, err
	}

	enumWindowsMutex.Lock()
	defer enumWindowsMutex.Unlock()

	enumWindowsList = make([]WindowInfo, 0, 100)
	enumProcessMap = pMap

	procEnumWindows.Call(enumWindowsCallback, 0)

	// Clean up map reference so GC can reclaim it
	enumProcessMap = nil

	// Copy to a new slice to return safely
	res := make([]WindowInfo, len(enumWindowsList))
	copy(res, enumWindowsList)
	return res, nil
}

func parseKeywordList(keywords string) []string {
	keywords = strings.ReplaceAll(keywords, "\r\n", "\n")
	parts := strings.FieldsFunc(keywords, func(r rune) bool {
		return r == ',' || r == '\n' || r == ';' || r == '|'
	})
	seen := map[string]bool{}
	var list []string
	for _, p := range parts {
		trimmed := strings.TrimSpace(strings.ToLower(p))
		trimmed = strings.TrimSuffix(trimmed, ".exe")
		if trimmed == "" || seen[trimmed] {
			continue
		}
		seen[trimmed] = true
		list = append(list, trimmed)
	}
	return list
}

var keywordAliases = map[string][]string{
	"lark":   {"lark", "feishu", "larkshell", "larkhelper"},
	"feishu": {"lark", "feishu", "larkshell", "larkhelper"},
	"teams":  {"teams", "ms-teams", "msteams"},
	"zalo":   {"zalo", "zalopcb"},
}

func expandKeyword(kw string) []string {
	base := []string{kw}
	if aliases, ok := keywordAliases[kw]; ok {
		base = append(base, aliases...)
	}
	seen := map[string]bool{}
	var out []string
	for _, a := range base {
		a = strings.TrimSpace(strings.ToLower(a))
		if a != "" && !seen[a] {
			seen[a] = true
			out = append(out, a)
		}
	}
	return out
}

func matchesAllowedKeyword(kw, pNameLower, wTitleLower string) bool {
	for _, variant := range expandKeyword(kw) {
		if strings.Contains(pNameLower, variant) || strings.Contains(wTitleLower, variant) {
			return true
		}
	}
	return false
}

func (b *Blocker) SetKeywords(keywords string) {
	b.mu.Lock()
	defer b.mu.Unlock()

	b.allowedKeywords = parseKeywordList(keywords)
	log.Printf("[BLOCKER] Keywords updated: %v", b.allowedKeywords)
}

var systemAllowed = map[string]bool{
	"explorer.exe":       true,
	"taskmgr.exe":        true,
	"cmd.exe":            true,
	"powershell.exe":     true,
	"conhost.exe":        true,
	"client.exe":         true, // App Wails của ta
	"wails.exe":          true,
	"msedgewebview2.exe": true, // WebView2 runtime của Wails
	"code.exe":           true, // VSCode
	"cursor.exe":         true, // Cursor
	"windsurf.exe":       true, // Windsurf
	"goland.exe":         true, // GoLand
	"goland64.exe":       true, // GoLand
	"idea64.exe":         true, // IntelliJ IDEA
	"clion64.exe":        true, // CLion
	"webstorm64.exe":     true, // WebStorm
	"pycharm64.exe":      true, // PyCharm
	"rider64.exe":        true, // Rider
	"studio64.exe":       true, // Android Studio
	"eclipse.exe":        true, // Eclipse
	"sublime_text.exe":   true, // Sublime Text
	"notepad++.exe":      true, // Notepad++
	"devenv.exe":         true, // Visual Studio
	"git-bash.exe":       true, // Git Bash
	"bash.exe":           true, // Bash
}

func (b *Blocker) checkAndKill() {
	b.mu.Lock()
	keywords := make([]string, len(b.allowedKeywords))
	copy(keywords, b.allowedKeywords)
	b.mu.Unlock()

	// Nếu không cấu hình keyword thì không chặn gì cả
	if len(keywords) == 0 {
		return
	}

	windows, err := EnumerateGUIWindows()
	if err != nil {
		return
	}

	for _, w := range windows {
		pNameLower := strings.ToLower(w.ProcessName)
		wTitleLower := strings.ToLower(w.Title)

		// 1. Luôn cho phép hệ thống/app cốt lõi hoặc chính tiến trình này
		if w.PID == uint32(os.Getpid()) || systemAllowed[pNameLower] || strings.Contains(pNameLower, "antigravity") || strings.Contains(wTitleLower, "antigravity") {
			continue
		}

		// 2. Kiểm tra xem có chứa bất kỳ từ khóa nào được cho phép không
		allowed := false
		for _, kw := range keywords {
			if matchesAllowedKeyword(kw, pNameLower, wTitleLower) {
				allowed = true
				break
			}
		}

		// 3. Nếu không nằm trong whitelist, tắt ứng dụng
		if !allowed {
			if b.OnBlocked != nil {
				b.OnBlocked(w.ProcessName, w.Title)
			}
			log.Printf("[BLOCKER] KILLED unauthorized application: %s (PID: %d, Title: %s)", w.ProcessName, w.PID, w.Title)
			h, err := syscall.OpenProcess(syscall.PROCESS_TERMINATE, false, w.PID)
			if err == nil {
				errTerm := syscall.TerminateProcess(h, 0)
				_ = syscall.CloseHandle(h)
				if errTerm == nil && b.OnKill != nil {
					b.OnKill(w.ProcessName, w.Title)
				}
			}
		}
	}
}

func (b *Blocker) Start() {
	b.mu.Lock()
	if b.running {
		b.mu.Unlock()
		return
	}
	b.running = true
	b.stopChan = make(chan struct{})
	b.mu.Unlock()

	go func() {
		ticker := time.NewTicker(3 * time.Second)
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				b.checkAndKill()
			case <-b.stopChan:
				return
			}
		}
	}()
	log.Println("[BLOCKER] Application Blocker Daemon Started.")
}

func (b *Blocker) Stop() {
	b.mu.Lock()
	defer b.mu.Unlock()

	if !b.running {
		return
	}
	b.running = false
	close(b.stopChan)
	log.Println("[BLOCKER] Application Blocker Daemon Stopped.")
}
