package main

import (
	"encoding/json"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// WebSocket upgrader配置
var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		// 允许所有源访问，生产环境中应限制具体域名
		return true
	},
}

// 客户端连接信息存储
type ClientInfo struct {
	DeviceType string
	RemoteAddr string
	Conn       *websocket.Conn
}

// 存储已连接的客户端信息
var (
	clients     = make(map[string]ClientInfo)
	clientsMutex = &sync.Mutex{}
)

// 保存图片数据到文件
func saveImageToFile(imageData []byte, deviceType string, clientAddr string) error {
	// 获取当前时间，用于生成唯一的文件名
	timestamp := time.Now().Format("20060102_150405_000")
	// 提取客户端IP地址的一部分作为标识符
	clientId := clientAddr
	if idx := strings.LastIndex(clientAddr, ":"); idx > 0 {
		clientId = clientAddr[:idx]
	}
	// 移除IP地址中的点号，使文件名更友好
	clientId = strings.ReplaceAll(clientId, ".", "_")
	// 构建文件名
	fileName := "screenshot_" + deviceType + "_" + clientId + "_" + timestamp + ".jpg"
	
	// 确保保存图片的目录存在
	imageDir := "images"
	if err := os.MkdirAll(imageDir, 0755); err != nil {
		return err
	}
	
	// 构建完整的文件路径
	filePath := filepath.Join(imageDir, fileName)
	
	// 写入文件
	if err := os.WriteFile(filePath, imageData, 0644); err != nil {
		return err
	}
	
	log.Printf("图片已保存到: %s，文件大小: %d 字节", filePath, len(imageData))
	return nil
}

// WebSocket连接处理函数
func wsHandler(w http.ResponseWriter, r *http.Request) {
	log.Printf("收到WebSocket连接请求，路径: %s", r.URL.Path)
	log.Printf("请求头 - Connection: %s, Upgrade: %s", 
		r.Header.Get("Connection"), r.Header.Get("Upgrade"))
	log.Printf("请求头 - Sec-WebSocket-Protocol: %s", r.Header.Get("Sec-WebSocket-Protocol"))

	// 从WebSocket子协议中获取设备信息
	device := "Unknown"
	protocols := r.Header.Get("Sec-WebSocket-Protocol")
	if protocols != "" {
		// 检查是否包含设备信息格式
		if protocols == "device-web" {
			// 直接匹配子协议名称为device-web的情况
			device = "Web"
		} else if protocols == "device-phone" {
			// 手机端的子协议（包括安卓和iPhone）
			device = "Phone"
		}
	}

	// 创建自定义header
	customHeaders := http.Header{}
	// customHeaders.Set("Device", "Web")

	// 创建临时upgrader，设置子协议以匹配客户端请求
	tempUpgrader := upgrader
	if protocols != "" {
		// 设置服务器支持的子协议，与客户端请求的子协议匹配
		log.Printf("接受客户端请求的子协议: %s", protocols)
		tempUpgrader.Subprotocols = []string{protocols}
	}

	// 将HTTP连接升级为WebSocket连接，并添加自定义header
	conn, err := tempUpgrader.Upgrade(w, r, customHeaders)
	if err != nil {
		log.Printf("升级WebSocket连接失败: %v", err)
		return
	}
	defer conn.Close()

	// 获取客户端地址
	clientAddr := conn.RemoteAddr().String()
	log.Printf("新的WebSocket连接建立，设备类型: %s, 本地地址: %s, 远程地址: %s", 
		device, conn.LocalAddr(), clientAddr)

	// 存储客户端信息
	clientsMutex.Lock()
	clients[clientAddr] = ClientInfo{
		DeviceType: device,
		RemoteAddr: clientAddr,
		Conn:       conn,
	}
	clientsMutex.Unlock()

	// 连接关闭时从列表中移除
	defer func() {
		clientsMutex.Lock()
		delete(clients, clientAddr)
		clientsMutex.Unlock()
	}()

	// 设置无超时连接，连接不会因为空闲而自动关闭
	conn.SetReadDeadline(time.Time{})

	// 发送欢迎消息给客户端，包含设备类型信息
	welcomeMsg := []byte("欢迎" + device + "连接到WebSocket服务器")
	log.Printf("尝试发送欢迎消息")
	if err := conn.WriteMessage(websocket.TextMessage, welcomeMsg); err != nil {
		log.Printf("发送欢迎消息失败: %v", err)
	} else {
		log.Printf("欢迎消息发送成功")
	}

	// 已移除心跳机制，使用无超时连接
	// 客户端通过onclose和onerror事件检测连接异常断开
	// 应用层消息确认机制和客户端重连策略由客户端实现

	// 循环读取客户端发送的消息
	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			log.Printf("读取消息失败: %v", err)
			break
		}

		// 处理二进制消息（图片数据）
		if messageType == websocket.BinaryMessage {
			log.Printf("收到二进制消息，长度: %d 字节，发送方设备类型: %s", len(p), device)
			
			// 如果是来自手机客户端的二进制消息，识别为图片数据并转发给所有Web客户端
			if device == "Phone" {
				// 保存图片到应用程序目录
				if err := saveImageToFile(p, device, clientAddr); err != nil {
					log.Printf("保存图片失败: %v", err)
				}
				
				clientsMutex.Lock()
				for _, client := range clients {
					if client.DeviceType == "Web" {
						// 向Web客户端转发二进制图片数据
						if err := client.Conn.WriteMessage(websocket.BinaryMessage, p); err != nil {
							log.Printf("向Web客户端 %s 转发图片数据失败: %v", client.RemoteAddr, err)
						} else {
							log.Printf("成功向Web客户端 %s 转发图片数据，长度: %d 字节", client.RemoteAddr, len(p))
						}
					}
				}
				clientsMutex.Unlock()
			}
			continue
		}

		messageStr := string(p)
		log.Printf("收到文本消息: %s", messageStr)

		// 检查消息格式是否为 messageId:content
		// 这种格式用于支持客户端的消息确认机制
		parts := strings.SplitN(messageStr, ":", 2)
		if len(parts) == 2 {
			// 提取消息ID和内容
			messageId := parts[0]
			content := parts[1]
			log.Printf("解析消息 - ID: %s, 内容: %s", messageId, content)

			// 发送确认消息给客户端
			confirmMsg := []byte("CONFIRM:" + messageId)
			if err := conn.WriteMessage(websocket.TextMessage, confirmMsg); err != nil {
				log.Printf("发送确认消息失败: %v", err)
			}

			// 检查是否为截图命令，如果是则转发给手机客户端
			if content == "screenshot" {
				log.Printf("检测到截图命令")
				// 获取手机客户端地址
				phoneAddr := getPhoneClientAddr()
				if phoneAddr != "" {
					// 向手机客户端发送截图命令
					clientsMutex.Lock()
					phoneClient, exists := clients[phoneAddr]
					clientsMutex.Unlock()
					
					if exists {
						// 直接向手机客户端发送截图命令
						log.Printf("正在向手机客户端 %s 发送截图命令", phoneAddr)
						if err := phoneClient.Conn.WriteMessage(websocket.TextMessage, []byte("screenshot")); err != nil {
							log.Printf("向手机客户端发送截图命令失败: %v", err)
							// 向Web客户端发送错误消息
							if err := conn.WriteMessage(websocket.TextMessage, []byte("向手机客户端发送截图命令失败")); err != nil {
								log.Printf("发送错误消息失败: %v", err)
							}
						} else {
							// 向Web客户端发送确认消息
							if err := conn.WriteMessage(websocket.TextMessage, []byte("截图命令已发送到手机客户端")); err != nil {
								log.Printf("发送确认消息失败: %v", err)
							}
						}
					} else {
						log.Printf("未找到手机客户端连接")
						if err := conn.WriteMessage(websocket.TextMessage, []byte("未找到手机客户端")); err != nil {
							log.Printf("发送错误消息失败: %v", err)
						}
					}
				} else {
					// 将原始内容返回给客户端
					if err := conn.WriteMessage(messageType, []byte(content)); err != nil {
						log.Printf("发送消息失败: %v", err)
						break
					}
				}
			} else {
				// 普通消息格式，直接返回
				log.Printf("收到普通格式消息")
				if err := conn.WriteMessage(messageType, p); err != nil {
					log.Printf("发送消息失败: %v", err)
					break
				}
			}
		}
	}
	log.Println("WebSocket连接已关闭")
}

// 判断是否为WebSocket升级请求
func isWebSocketRequest(r *http.Request) bool {
	return r.Header.Get("Connection") == "Upgrade" && r.Header.Get("Upgrade") == "websocket"
}

// 根处理器，根据请求类型分发
func rootHandler(w http.ResponseWriter, r *http.Request) {
	if isWebSocketRequest(r) {
		wsHandler(w, r)
		return
	}
	httpHandler(w, r)
}

// 获取手机客户端地址
func getPhoneClientAddr() string {
	clientsMutex.Lock()
	defer clientsMutex.Unlock()
	// 注释掉不需要的日志
	// log.Printf("当前连接的客户端数量: %d", len(clients))
	var phoneAddr string
	for _, client := range clients {
		// log.Printf("客户端 [%s] - 设备类型: %s, 远程地址: %s", addr, client.DeviceType, client.RemoteAddr)
		if client.DeviceType == "Phone" {
			// log.Printf("找到手机客户端，地址: %s", addr)
			phoneAddr = client.RemoteAddr
			break // 找到一个手机客户端就返回
		}
	}
	// log.Println("未找到手机客户端")
	return phoneAddr
}

// HTTP请求处理函数
func httpHandler(w http.ResponseWriter, r *http.Request) {
	// 获取请求路径
	path := r.URL.Path

	// API请求处理
	if strings.HasPrefix(path, "/api/") {
		handleAPIRequest(w, r, path)
		return
	}

	// 详细记录HTTP请求信息
	log.Printf("接收到HTTP请求 - 路径: %s, 方法: %s", path, r.Method)

	// SPA应用路由处理优化：对于非文件路径（没有扩展名），直接返回index.html
	// 这确保了所有前端路由路径都能被正确处理
	if !strings.Contains(path, ".") && path != "/" {
		log.Printf("检测到SPA路由路径: %s，直接返回index.html", path)
		handleSPA(w, r)
		return
	}

	// 静态文件服务 - 优先从dist目录提供构建后的文件
	distHandler := http.FileServer(http.Dir("./Web/dist"))
	distReq, err := http.NewRequest(r.Method, path, r.Body)
	if err != nil {
		// 如果创建请求失败，尝试从Web目录提供文件
		log.Printf("创建dist请求失败: %v，尝试从Web目录提供文件", err)
		handleStaticFileFromWeb(w, r, path)
		return
	}

	// 复制原始请求的所有头信息
	*distReq = *r
	
	// 尝试在dist目录下查找文件
	recorder := httptest.NewRecorder()
	distHandler.ServeHTTP(recorder, distReq)
	
	// 检查是否找到了文件
	if recorder.Code != http.StatusNotFound {
		log.Printf("在dist目录找到文件，返回状态码: %d", recorder.Code)
		// 找到了文件，将响应复制回客户端
		for k, v := range recorder.Header() {
			w.Header()[k] = v
		}
		w.WriteHeader(recorder.Code)
		w.Write(recorder.Body.Bytes())
		return
	}

	log.Printf("在dist目录未找到文件，尝试从Web目录提供文件")
	// dist目录下未找到文件，尝试从Web目录提供文件
	handleStaticFileFromWeb(w, r, path)
}

// 处理SPA路由请求
func handleSPA(w http.ResponseWriter, r *http.Request) {
	log.Printf("执行SPA路由处理 - 请求路径: %s", r.URL.Path)
	
	// 获取当前工作目录
	cwd, _ := os.Getwd()
	distDir := filepath.Join(cwd, "Web", "dist")
	
	// 检查dist目录是否存在
	if _, err := os.Stat(distDir); os.IsNotExist(err) {
		log.Printf("警告: dist目录不存在: %s，尝试从Web目录提供index.html", distDir)
		// 如果dist目录不存在，尝试从Web目录提供index.html
		handleStaticFileFromWeb(w, r, "/index.html")
		return
	}
	
	// 从dist目录读取index.html
	indexPath := filepath.Join(distDir, "index.html")
	content, err := os.ReadFile(indexPath)
	if err != nil {
		log.Printf("读取index.html失败: %v", err)
		// 如果读取失败，尝试从Web目录提供index.html
		handleStaticFileFromWeb(w, r, "/index.html")
		return
	}
	
	// 设置响应头
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	// 确保CORS设置允许所有来源
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers", "*")
	// 写入index.html内容
	w.WriteHeader(http.StatusOK)
	w.Write(content)
	log.Printf("SPA路由处理成功，返回index.html")
}

// 处理API请求
func handleAPIRequest(w http.ResponseWriter, r *http.Request, path string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	// 处理获取手机客户端地址的API请求
	if path == "/api/android-addr" {
		// 保持向后兼容
		phoneAddr := getPhoneClientAddr()
		w.Write([]byte(`{"status":"ok","androidAddress":"` + phoneAddr + `"}`))
		return
	}
	
	// 统一的手机客户端地址API
	if path == "/api/phone-addr" {
		phoneAddr := getPhoneClientAddr()
		w.Write([]byte(`{"status":"ok","phoneAddress":"` + phoneAddr + `"}`))
		return
	}

	// 其他API请求处理
	response := map[string]interface{}{
		"status": "ok",
		"message": "This is an API response",
		"path": path,
		"timestamp": time.Now().Unix(),
	}

	// 将响应转换为JSON
	jsonData, err := json.Marshal(response)
	if err != nil {
		log.Printf("JSON序列化失败: %v", err)
		w.Write([]byte(`{"status":"error","message":"Internal server error"}`))
		return
	}

	w.Write(jsonData)
}

// 从Web目录提供静态文件
func handleStaticFileFromWeb(w http.ResponseWriter, r *http.Request, path string) {
	// 详细记录请求路径，帮助诊断SPA路由问题
	log.Printf("处理静态文件请求 - 路径: %s, 方法: %s", path, r.Method)

	// 如果请求根路径，则返回index.html
	if path == "/" {
		log.Printf("请求根路径，返回index.html")
		path = "/index.html"
	}

	// 获取当前工作目录，确保路径正确
	cwd, _ := os.Getwd()
	log.Printf("当前工作目录: %s", cwd)

	// 检查Web目录是否存在
	webDir := filepath.Join(cwd, "Web")
	if _, err := os.Stat(webDir); os.IsNotExist(err) {
		log.Printf("警告: Web目录不存在: %s", webDir)
	}

	webHandler := http.FileServer(http.Dir(webDir))
	webReq, err := http.NewRequest(r.Method, path, r.Body)
	if err != nil {
		log.Printf("创建Web请求失败: %v", err)
		// 在错误情况下也返回index.html，确保SPA路由正常工作
		path = "/index.html"
		webReq, _ = http.NewRequest(r.Method, path, r.Body)
	}

	// 复制原始请求的所有头信息
	*webReq = *r
	
	// 尝试在Web目录下查找文件
	recorder := httptest.NewRecorder()
	webHandler.ServeHTTP(recorder, webReq)
	
	// 检查是否找到了文件
	if recorder.Code != http.StatusNotFound {
		log.Printf("文件找到，返回状态码: %d", recorder.Code)
		// 找到了文件，将响应复制回客户端
		for k, v := range recorder.Header() {
			w.Header()[k] = v
		}
		w.WriteHeader(recorder.Code)
		w.Write(recorder.Body.Bytes())
		return
	}

	// SPA应用处理：当找不到具体文件时，返回index.html，让前端路由处理
	log.Printf("文件未找到，启动SPA路由处理 - 路径: %s", path)
	indexReq, _ := http.NewRequest(r.Method, "/index.html", r.Body)
	*indexReq = *r
	indexRecorder := httptest.NewRecorder()
	webHandler.ServeHTTP(indexRecorder, indexReq)
	
	log.Printf("SPA路由处理 - index.html返回状态码: %d", indexRecorder.Code)
	for k, v := range indexRecorder.Header() {
		w.Header()[k] = v
	}
	w.WriteHeader(indexRecorder.Code)
	w.Write(indexRecorder.Body.Bytes())
}

// 将未找到的请求作为API响应处理
func handleNotFoundAsAPI(w http.ResponseWriter, r *http.Request, path string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	response := map[string]interface{}{
		"status": "not_found",
		"message": "Resource not found",
		"path": path,
		"timestamp": time.Now().Unix(),
	}

	// 将响应转换为JSON
	jsonData, err := json.Marshal(response)
	if err != nil {
		log.Printf("JSON序列化失败: %v", err)
		w.Write([]byte(`{"status":"error","message":"Internal server error"}`))
		return
	}

	w.Write(jsonData)
}

func main() {
	// 根路径处理器，根据请求类型分发到WebSocket或HTTP处理器
	http.HandleFunc("/", rootHandler)

	// 启动服务器
	port := ":8805"
	log.Printf("服务器已启动，监听端口 %s...", port)
	log.Printf("前端页面: http://localhost%s", port)
	log.Printf("WebSocket连接: ws://localhost%s", port)

	if err := http.ListenAndServe(port, nil); err != nil {
		log.Fatalf("服务器启动失败: %v", err)
	}
}