package main

import (
	"encoding/json"
	"fmt"
	"math"
	"os"
	"strings"
)

func loadAnnotation(path string) (AnnotationFile, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return AnnotationFile{}, err
	}

	var annotation AnnotationFile
	if err := json.Unmarshal(data, &annotation); err != nil {
		return AnnotationFile{}, err
	}

	return annotation, nil
}

func summarizeRegions(regions []Region) []RegionSummary {
	summaries := make([]RegionSummary, 0, len(regions))
	for _, region := range regions {
		bboxSource := region.BBoxSource
		if isZeroBox(bboxSource) {
			scaleFactor := region.ScaleFactor
			if scaleFactor == 0 {
				scaleFactor = 1
			}
			bboxSource = [4]float64{
				region.BBoxImage[0] / scaleFactor,
				region.BBoxImage[1] / scaleFactor,
				region.BBoxImage[2] / scaleFactor,
				region.BBoxImage[3] / scaleFactor,
			}
		}

		summaries = append(summaries, RegionSummary{
			Name:         region.Name,
			Description:  region.Description,
			Type:         region.Type,
			BBoxScreen:   region.BBoxScreen,
			BBoxSource:   bboxSource,
			CenterScreen: centerOf(region.BBoxScreen),
			CenterSource: centerOf(bboxSource),
		})
	}
	return summaries
}

func centerOf(box [4]float64) [2]float64 {
	return [2]float64{(box[0] + box[2]) / 2, (box[1] + box[3]) / 2}
}

func regionByType(summaries []RegionSummary) map[string]RegionSummary {
	index := make(map[string]RegionSummary, len(summaries))
	for _, summary := range summaries {
		if _, exists := index[summary.Type]; !exists {
			index[summary.Type] = summary
		}
	}
	return index
}

func validateRequiredRegions(summaries []RegionSummary) error {
	index := regionByType(summaries)
	missing := make([]string, 0)
	for _, regionType := range []string{"chat_content", "input_box"} {
		if _, exists := index[regionType]; !exists {
			missing = append(missing, regionType)
		}
	}
	if len(missing) > 0 {
		return fmt.Errorf("missing required regions: %v", missing)
	}
	return nil
}

func inferMissingRegionTypes(summaries []RegionSummary) ([]RegionSummary, []string) {
	index := regionByType(summaries)
	warnings := make([]string, 0)
	customs := make([]RegionSummary, 0)
	for _, summary := range summaries {
		if inferredType := inferTypeFromDescription(summary.Description); inferredType != "" {
			if _, exists := index[inferredType]; !exists {
				inferred := summary
				inferred.Type = inferredType
				inferred.Name = inferred.Name + "（按描述推断）"
				summaries = append(summaries, inferred)
				index[inferredType] = inferred
				warnings = append(warnings, fmt.Sprintf("区域 %q 已按描述推断为 %s", summary.Name, inferredType))
			}
		}

		if summary.Type == "custom" {
			customs = append(customs, summary)
		}
	}
	if len(customs) == 0 {
		return summaries, warnings
	}

	if _, exists := index["contact_list"]; !exists {
		best := customs[0]
		for _, candidate := range customs[1:] {
			if candidate.CenterScreen[0] < best.CenterScreen[0] && boxHeight(candidate.BBoxScreen) > boxHeight(best.BBoxScreen)*0.5 {
				best = candidate
			}
		}
		best.Type = "contact_list"
		best.Name = best.Name + "（推断联系人区域）"
		summaries = append(summaries, best)
		index[best.Type] = best
		warnings = append(warnings, "未标注 contact_list，已按最左侧大区域推断")
	}

	if _, exists := index["input_box"]; !exists {
		best := customs[0]
		for _, candidate := range customs[1:] {
			if candidate.CenterScreen[1] > best.CenterScreen[1] || nearlyEqual(candidate.CenterScreen[1], best.CenterScreen[1]) && boxArea(candidate.BBoxScreen) > boxArea(best.BBoxScreen) {
				best = candidate
			}
		}
		best.Type = "input_box"
		best.Name = best.Name + "（推断输入框区域）"
		summaries = append(summaries, best)
		index[best.Type] = best
		warnings = append(warnings, "未标注 input_box，已按底部区域推断")
	}

	if _, exists := index["chat_content"]; !exists {
		input := index["input_box"]
		contact := index["contact_list"]
		best := customs[0]
		bestScore := math.Inf(-1)
		for _, candidate := range customs {
			if sameBox(candidate.BBoxScreen, input.BBoxScreen) || sameBox(candidate.BBoxScreen, contact.BBoxScreen) {
				continue
			}
			score := boxArea(candidate.BBoxScreen) - math.Abs(candidate.CenterScreen[1]-input.CenterScreen[1])*0.2
			if score > bestScore {
				best = candidate
				bestScore = score
			}
		}
		best.Type = "chat_content"
		best.Name = best.Name + "（推断聊天内容区域）"
		summaries = append(summaries, best)
		index[best.Type] = best
		warnings = append(warnings, "未标注 chat_content，已按大面积内容区推断")
	}

	return summaries, warnings
}

func sameBox(a, b [4]float64) bool {
	return nearlyEqual(a[0], b[0]) && nearlyEqual(a[1], b[1]) && nearlyEqual(a[2], b[2]) && nearlyEqual(a[3], b[3])
}

func isZeroBox(box [4]float64) bool {
	return box[0] == 0 && box[1] == 0 && box[2] == 0 && box[3] == 0
}

func boxWidth(box [4]float64) float64 {
	return math.Abs(box[2] - box[0])
}

func boxHeight(box [4]float64) float64 {
	return math.Abs(box[3] - box[1])
}

func boxArea(box [4]float64) float64 {
	return boxWidth(box) * boxHeight(box)
}

func nearlyEqual(a, b float64) bool {
	return math.Abs(a-b) < 8
}

func inferTypeFromDescription(description string) string {
	text := strings.ToLower(description)
	if text == "" {
		return ""
	}

	if strings.Contains(text, "chat_content") || strings.Contains(text, "聊天") || strings.Contains(text, "消息") || strings.Contains(text, "记录") {
		return "chat_content"
	}
	if strings.Contains(text, "input_box") || strings.Contains(text, "输入") || strings.Contains(text, "文本框") || strings.Contains(text, "编辑框") {
		return "input_box"
	}
	if strings.Contains(text, "send_button") || strings.Contains(text, "发送") || strings.Contains(text, "send") {
		return "send_button"
	}
	if strings.Contains(text, "contact_list") || strings.Contains(text, "联系人") || strings.Contains(text, "会话列表") || strings.Contains(text, "好友") {
		return "contact_list"
	}

	return ""
}