package server

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"vdi-cloud/platform/internal/model"
	"vdi-cloud/platform/internal/store"
)

func (s *Server) knownGPURecovery(state gpuPrepareState) bool {
	return state.UPID != "" && state.Node == s.Config.ExclusiveGPUNode &&
		fmt.Sprint(state.Members) == fmt.Sprint(s.Config.ExclusiveGPUVMIDs) &&
		(state.Phase == "start-waiting" || state.Phase == "shutdown-waiting" || state.Phase == "service-waiting")
}

func (s *Server) gpuRecoveryPending(task model.Task) bool {
	var state gpuPrepareState
	return task.Status == "timeout" && time.Since(task.CreatedAt) < 5*time.Minute &&
		json.Unmarshal([]byte(task.ProviderRef), &state) == nil && s.knownGPURecovery(state)
}

// Recovery observes the original UPID and every member under the same group
// lock as preparation. It never sends a power command, even after an API
// restart or a lost submission response.
func (s *Server) recoverGPUPrepare(parent context.Context, task model.Task) {
	ctx, cancel := context.WithTimeout(parent, 60*time.Second)
	defer cancel()
	if s.RefreshCompute(ctx) != nil {
		return
	}
	pve := s.currentPVE()
	if pve == nil {
		return
	}
	tx, err := s.DB.BeginTxx(ctx, nil)
	if err != nil {
		return
	}
	defer tx.Rollback()
	locked, err := tryGPULock(ctx, tx)
	if err != nil || !locked {
		return
	}
	if err = tx.GetContext(ctx, &task, "SELECT * FROM tasks WHERE id=$1 AND kind='gpu-prepare' AND status IN ('running','timeout') FOR UPDATE", task.ID); err != nil {
		return
	}
	finish := func(status, message string, power map[string]string) {
		for id, value := range power {
			if _, err = tx.ExecContext(ctx, "UPDATE desktops SET power_state=$1,service_state=CASE WHEN $2=$3 AND $4='success' THEN 'ready' ELSE 'unknown' END,updated_at=now() WHERE id=$2", value, id, task.TargetID, status); err != nil {
				return
			}
		}
		if _, err = tx.ExecContext(ctx, "UPDATE tasks SET status=$1,message=$2,updated_at=now() WHERE id=$3", status, message, task.ID); err != nil {
			return
		}
		if status != "timeout" {
			if _, err = tx.ExecContext(ctx, "INSERT INTO audit_logs VALUES($1,'system','auto-reconcile-gpu-switch',$2,$3,$4,now())", store.ID("audit_"), task.TargetID, status, task.ID); err != nil {
				return
			}
		}
		_ = tx.Commit()
	}
	var state gpuPrepareState
	if json.Unmarshal([]byte(task.ProviderRef), &state) == nil && state.Phase == "queued" && state.UPID == "" {
		finish("failed", "服务重启前尚未提交电源操作，可重新选择桌面", nil)
		return
	}
	if json.Unmarshal([]byte(task.ProviderRef), &state) != nil || !s.knownGPURecovery(state) {
		finish("timeout", "未记录可核对的原任务或独占组配置已变化，请管理员核查；不会重复发送电源命令", nil)
		return
	}
	original, err := pve.Task(ctx, state.Node, state.UPID)
	if err != nil || original.Status != "stopped" || original.ExitStatus == "" {
		finish("timeout", "平台正在自动核对原 PVE 任务，暂不重复切换；计算节点恢复后将继续检查", nil)
		return
	}
	members, err := s.gpuMembers(ctx, tx)
	if err != nil {
		finish("timeout", "独占显卡组成员不完整，请管理员核查", nil)
		return
	}
	powers := map[string]string{}
	running, matches, submittedMember := 0, true, false
	var target model.Desktop
	for _, d := range members {
		power, readErr := pve.Power(ctx, d.Node, d.VMID)
		if readErr != nil || (power != "running" && power != "stopped") {
			finish("timeout", "平台正在自动核对虚拟机实际电源状态，暂不重复切换", nil)
			return
		}
		powers[d.ID] = power
		if d.VMID == state.VMID {
			submittedMember = true
		}
		if power == "running" {
			running++
		}
		if d.ID == task.TargetID {
			target = d
			matches = matches && power == "running"
		} else {
			matches = matches && power == "stopped"
		}
	}
	if running > 1 || target.ID == "" || !submittedMember {
		finish("timeout", "实际电源或显卡归属存在冲突，请管理员核查；平台不会自动关机", powers)
		return
	}
	if original.ExitStatus != "OK" || !matches {
		finish("failed", "已自动核对原任务和电源状态，本次准备未完成；可重新选择桌面", powers)
		return
	}
	var allowed bool
	err = tx.GetContext(ctx, &allowed, `SELECT EXISTS(SELECT 1 FROM bindings b JOIN users u ON u.id=b.user_id JOIN terminals t ON t.id=u.terminal_id WHERE b.desktop_id=$1 AND t.id=$2 AND u.enabled AND t.enabled)`, target.ID, state.Terminal)
	if err != nil {
		return
	}
	if !allowed || target.SetupRequired || target.Endpoint == "" {
		finish("failed", "原任务已核对，但桌面授权或连接配置已变化，请检查配置后重新选择", powers)
		return
	}
	if !sunshineReady(ctx, target.Endpoint) {
		if time.Since(task.CreatedAt) < 5*time.Minute {
			finish("timeout", "原任务已确认成功，正在等待桌面连接服务恢复", powers)
		} else {
			finish("failed", "原任务及显卡归属已自动确认，但 Sunshine 仍未响应，请检查桌面服务", powers)
		}
		return
	}
	finish("success", "已自动核对原任务、独占显卡归属和连接服务，桌面可连接", powers)
}

// Only called by the worker while no GPU/provisioning operation is running.
// A restart resumes observations, never the command sequence.
func (s *Server) nextGPURecovery(ctx context.Context) (model.Task, error) {
	var task model.Task
	err := s.DB.GetContext(ctx, &task, `SELECT * FROM tasks WHERE kind='gpu-prepare'
AND ((status='running')
 OR (status='timeout' AND provider_ref::jsonb->>'upid' IS NOT NULL AND provider_ref::jsonb->>'upid' != '' AND created_at>now()-interval '24 hours'))
ORDER BY updated_at,id LIMIT 1`)
	return task, err
}