package main

import (
	"encoding/json"
	"fmt"
	"sort"
	"strings"
)

type AXProbeResult struct {
	AppName string   `json:"app_name"`
	PID     int      `json:"pid"`
	Nodes   []AXNode `json:"nodes"`
}

type AXNode struct {
	Role        string     `json:"role"`
	Title       string     `json:"title"`
	Value       string     `json:"value"`
	Description string     `json:"description"`
	Frame       [4]float64 `json:"frame"`
	Depth       int        `json:"depth"`
}

func runAXProbe() error {
	logInfo("macOS Accessibility 微信 UI 树探测启动")

	config, err := loadConfig("config.toml")
	if err != nil {
		return fmt.Errorf("读取 config.toml 失败：%w", err)
	}

	regionsPath := appDataPath(config, config.Agent.RegionsPath)
	annotation, err := loadAnnotation(regionsPath)
	if err != nil {
		return fmt.Errorf("读取标注文件失败：%w", err)
	}
	summaries, warnings := inferMissingRegionTypes(summarizeRegions(annotation.Regions))
	for _, warning := range warnings {
		logWarning(warning)
	}

	probe, err := probeWeChatAXTree()
	if err != nil {
		return err
	}
	logInfo(fmt.Sprintf("AX 节点总数=%d pid=%d app=%s", len(probe.Nodes), probe.PID, probe.AppName))

	textNodes := usefulAXTextNodes(probe.Nodes)
	logInfo(fmt.Sprintf("可读文本/按钮节点=%d", len(textNodes)))
	for index, node := range firstNodes(textNodes, 80) {
		logJSON("ax", fmt.Sprintf("#%02d depth=%d role=%s frame=[%.0f,%.0f,%.0f,%.0f] text=%s", index+1, node.Depth, node.Role, node.Frame[0], node.Frame[1], node.Frame[2], node.Frame[3], axNodeText(node)))
	}

	if chat, ok := regionByType(summaries)["chat_content"]; ok {
		chatNodes := nodesInsideBox(textNodes, chat.BBoxScreen)
		sort.SliceStable(chatNodes, func(i, j int) bool {
			if chatNodes[i].Frame[1] == chatNodes[j].Frame[1] {
				return chatNodes[i].Frame[0] < chatNodes[j].Frame[0]
			}
			return chatNodes[i].Frame[1] < chatNodes[j].Frame[1]
		})
		logInfo(fmt.Sprintf("chat_content 区域内可读节点=%d", len(chatNodes)))
		for index, node := range firstNodes(chatNodes, 80) {
			logJSON("ax-chat", fmt.Sprintf("#%02d role=%s frame=[%.0f,%.0f,%.0f,%.0f] text=%s", index+1, node.Role, node.Frame[0], node.Frame[1], node.Frame[2], node.Frame[3], axNodeText(node)))
		}
	} else {
		logWarning("未找到 chat_content 标注区域，跳过区域过滤")
	}

	return nil
}

func probeWeChatAXTree() (AXProbeResult, error) {
	script := `
import AppKit
import ApplicationServices
import Foundation

struct Node: Encodable {
  let role: String
  let title: String
  let value: String
  let description: String
  let frame: [Double]
  let depth: Int
}

struct Result: Encodable {
  let app_name: String
  let pid: Int
  let nodes: [Node]
}

func stringAttr(_ element: AXUIElement, _ key: String) -> String {
  var value: CFTypeRef?
  let error = AXUIElementCopyAttributeValue(element, key as CFString, &value)
  if error != .success || value == nil { return "" }
  if let str = value as? String { return clean(str) }
  if let attr = value as? NSAttributedString { return clean(attr.string) }
  return ""
}

func clean(_ value: String) -> String {
  return value
    .replacingOccurrences(of: "\n", with: " ")
    .replacingOccurrences(of: "\r", with: " ")
    .replacingOccurrences(of: "\t", with: " ")
}

func frameOf(_ element: AXUIElement) -> [Double] {
  var positionValue: CFTypeRef?
  var sizeValue: CFTypeRef?
  var point = CGPoint.zero
  var size = CGSize.zero
  if AXUIElementCopyAttributeValue(element, kAXPositionAttribute as CFString, &positionValue) == .success,
     let positionValue,
     CFGetTypeID(positionValue) == AXValueGetTypeID() {
    AXValueGetValue(positionValue as! AXValue, .cgPoint, &point)
  }
  if AXUIElementCopyAttributeValue(element, kAXSizeAttribute as CFString, &sizeValue) == .success,
     let sizeValue,
     CFGetTypeID(sizeValue) == AXValueGetTypeID() {
    AXValueGetValue(sizeValue as! AXValue, .cgSize, &size)
  }
  return [Double(point.x), Double(point.y), Double(size.width), Double(size.height)]
}

func childrenOf(_ element: AXUIElement) -> [AXUIElement] {
  var value: CFTypeRef?
  let error = AXUIElementCopyAttributeValue(element, kAXChildrenAttribute as CFString, &value)
  if error != .success || value == nil { return [] }
  return value as? [AXUIElement] ?? []
}

func walk(_ element: AXUIElement, depth: Int, nodes: inout [Node], visited: inout Int) {
  if visited > 2500 || depth > 20 { return }
  visited += 1
  let node = Node(
    role: stringAttr(element, kAXRoleAttribute),
    title: stringAttr(element, kAXTitleAttribute),
    value: stringAttr(element, kAXValueAttribute),
    description: stringAttr(element, kAXDescriptionAttribute),
    frame: frameOf(element),
    depth: depth
  )
  nodes.append(node)
  for child in childrenOf(element) {
    walk(child, depth: depth + 1, nodes: &nodes, visited: &visited)
  }
}

let trusted = AXIsProcessTrustedWithOptions([kAXTrustedCheckOptionPrompt.takeUnretainedValue() as String: true] as CFDictionary)
if !trusted {
  fputs("Accessibility permission is not granted\n", stderr)
  exit(10)
}

let apps = NSWorkspace.shared.runningApplications.filter { app in
  let name = app.localizedName ?? ""
  return name == "微信" || name == "WeChat" || name.lowercased().contains("wechat")
}

guard let app = apps.first else {
  fputs("WeChat app not found\n", stderr)
  exit(2)
}

app.activate(options: [.activateIgnoringOtherApps])
usleep(300000)

let root = AXUIElementCreateApplication(app.processIdentifier)
var nodes: [Node] = []
var visited = 0
walk(root, depth: 0, nodes: &nodes, visited: &visited)
let result = Result(app_name: app.localizedName ?? "", pid: Int(app.processIdentifier), nodes: nodes)
let data = try JSONEncoder().encode(result)
print(String(data: data, encoding: .utf8)!)
`

	output, err := runSwift(script)
	if err != nil {
		return AXProbeResult{}, fmt.Errorf("AX probe failed: %w", err)
	}

	var result AXProbeResult
	if err := json.Unmarshal([]byte(output), &result); err != nil {
		preview := output
		if len(preview) > 300 {
			preview = preview[:300]
		}
		return AXProbeResult{}, fmt.Errorf("parse AX JSON failed: %w output_prefix=%q", err, preview)
	}
	return result, nil
}

func usefulAXTextNodes(nodes []AXNode) []AXNode {
	result := make([]AXNode, 0)
	seen := map[string]bool{}
	for _, node := range nodes {
		text := strings.TrimSpace(axNodeText(node))
		if text == "" {
			continue
		}
		if node.Frame[2] <= 0 || node.Frame[3] <= 0 {
			continue
		}
		key := fmt.Sprintf("%s|%s|%.0f|%.0f|%.0f|%.0f", node.Role, text, node.Frame[0], node.Frame[1], node.Frame[2], node.Frame[3])
		if seen[key] {
			continue
		}
		seen[key] = true
		result = append(result, node)
	}
	sort.SliceStable(result, func(i, j int) bool {
		if result[i].Frame[1] == result[j].Frame[1] {
			return result[i].Frame[0] < result[j].Frame[0]
		}
		return result[i].Frame[1] < result[j].Frame[1]
	})
	return result
}

func axNodeText(node AXNode) string {
	parts := []string{node.Value, node.Title, node.Description}
	for _, part := range parts {
		part = strings.TrimSpace(part)
		if part != "" {
			return part
		}
	}
	return ""
}

func nodesInsideBox(nodes []AXNode, box [4]float64) []AXNode {
	result := make([]AXNode, 0)
	for _, node := range nodes {
		cx := node.Frame[0] + node.Frame[2]/2
		cy := node.Frame[1] + node.Frame[3]/2
		if cx >= box[0] && cx <= box[2] && cy >= box[1] && cy <= box[3] {
			result = append(result, node)
		}
	}
	return result
}

func firstNodes(nodes []AXNode, limit int) []AXNode {
	if len(nodes) <= limit {
		return nodes
	}
	return nodes[:limit]
}