package controllers

import (
	"aon_app_server/models"
	"aon_app_server/utils/mongo"
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/beego/beego/v2/core/logs"
	beego "github.com/beego/beego/v2/server/web"
	"go.mongodb.org/mongo-driver/bson"
	"gopkg.in/yaml.v2"
	"io"
	"net/http"
	"os"
	"reflect"
	"strings"
	"time"
)

type TaskController struct {
	MainController
}

var supportModels map[string]*models.Model

func init() {
	err := readYAML("./conf/replicate_models_version.yaml", &supportModels)
	if err != nil {
		logs.Debug("Error reading YAML file: %v", err)
	}
}

var execTasks = make(chan *models.Task, 100)

func (server *TaskController) Prediction() {
	body := server.Ctx.Input.RequestBody
	task := models.Task{}
	err := json.Unmarshal(body, &task) //解析body中数据
	logs.Debug("appRequest", task)
	if err != nil {
		server.respond(models.NoRequestBody, err.Error())
		return
	}

	asyncString := server.Ctx.Input.Header("Async")
	async := false
	if asyncString == "true" {
		async = true
	}
	task.CreatedTime = time.Now().UTC()
	task.UpdatedTime = task.CreatedTime
	id, err := mongo.Insert(&task)
	if err != nil {
		server.respond(models.BusinessFailed, err.Error())
		return
	}
	task.Id = id
	whois, _ := beego.AppConfig.String("whoisApi")
	var result *models.TaskResponse
	if whois == "aonet" {
		result, err = sendTask(&task, async)
	} else if whois == "replicate" {
		result, err = sendReplicate(&task, async)
	}
	if err != nil {
		server.respond(http.StatusOK, err.Error())
		return
	}
	server.respond(http.StatusOK, "", result)
}

func copyImages(images []string) []string {
	imagesToCopy := models.ImagesToCopy{
		Sources: images,
	}
	url, _ := beego.AppConfig.String("imageUrl")
	payload := new(bytes.Buffer)
	json.NewEncoder(payload).Encode(imagesToCopy)
	resp, err := http.Post(url, "application/json;charset=UTF-8", payload)
	if err != nil {
		logs.Info("Error sending request:", err)
		return images
	}
	defer resp.Body.Close()
	logs.Info("copyImages resp code", resp.StatusCode)
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		logs.Info("Error reading response:", err)
		return images
	}
	logs.Debug("copyImages body", string(body))

	var response []models.ImageCopied
	if err := json.Unmarshal(body, &response); err != nil {
		return images
	}
	var backImages []string
	for _, value := range response {
		if value.Code == 200 {
			backImages = append(backImages, value.Data)
		}
	}
	if len(backImages) > 0 {
		return backImages
	}
	return images
}

func transferImages(images []string) []string {
	imagesToCopy := models.ImagesToCopy{
		Sources: images,
	}
	url, _ := beego.AppConfig.String("imageTransferUrl")
	payload := new(bytes.Buffer)
	json.NewEncoder(payload).Encode(imagesToCopy)
	resp, err := http.Post(url, "application/json;charset=UTF-8", payload)
	if err != nil {
		logs.Info("Error sending request:", err)
		return images
	}
	defer resp.Body.Close()
	logs.Info("transferImages resp code", resp.StatusCode)
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		logs.Info("Error reading response:", err)
		return images
	}
	logs.Debug("transferImages body", string(body))

	var response []models.ImageCopied
	if err := json.Unmarshal(body, &response); err != nil {
		return images
	}
	var backImages []string
	for _, value := range response {
		if value.Code == 200 {
			backImages = append(backImages, value.Data)
		}
	}
	if len(backImages) > 0 {
		return backImages
	}
	return images
}

func sendTask(task *models.Task, async bool) (*models.TaskResponse, error) {
	host, _ := beego.AppConfig.String("taskUrl")
	url := host + task.ApiPath
	payload := new(bytes.Buffer)
	//var input interface{}
	//if err := json.Unmarshal([]byte(task.Input), &input); err != nil {
	//	setError(task, "task.Input Unmarshal error:"+err.Error())
	//	logs.Info("sendTask task.Input Unmarshal response:", err)
	//	return nil, err
	//}
	json.NewEncoder(payload).Encode(task.Input)
	client := &http.Client{}
	request, err := http.NewRequest("POST", url, payload)
	if err != nil {
		setError(task, "sendTask request create error:"+err.Error())
		logs.Info("sendTask Error NewRequest request:", err)
		return nil, err
	}
	apikey, _ := beego.AppConfig.String("apikey")
	request.Header.Add("apikey", apikey)

	logs.Info("sendTask client sending request:")
	resp, err := client.Do(request)
	if err != nil {
		setError(task, "Task sending error:"+err.Error())
		logs.Info("sendTask Error sending request:", err)
		return nil, err
	}
	defer resp.Body.Close()
	logs.Info("sendTask resp code", resp.StatusCode)
	body, err := io.ReadAll(resp.Body)
	logs.Info("sendTask body", string(body))
	if resp.StatusCode == 200 && body == nil {
		setError(task, "Task response body null")
		logs.Info("sendTask Body reading response:", err)
		return nil, err
	}
	if err != nil {
		setError(task, "Task read response body error:"+err.Error())
		logs.Info("sendTask Error reading response:", err)
		return nil, err
	}
	var response models.TaskResponse
	if err = json.Unmarshal(body, &response); err != nil {
		setError(task, "Task response Unmarshal error:"+err.Error())
		logs.Info("sendTask Error Unmarshal response:", err)
		return nil, err
	}
	taskId := resp.Header.Get("task-id")
	task.TaskId = taskId
	task.UpdatedTime = time.Now().UTC()
	if response.Task.IsSuccess {
		task.Status = 2
		task.Output = copyImages(response.Output)
		mongo.Update(task)
		response.Output = task.Output
	} else {
		task.Status = 3
		task.Error = response.Task
		_, err = mongo.Update(task)
		if err != nil {
			logs.Info("Update Task Error:", err)
		}
	}
	return &response, nil
}

func readYAML(filename string, out interface{}) error {
	data, err := os.ReadFile(filename)
	if err != nil {
		return err
	}

	err = yaml.Unmarshal(data, out)
	if err != nil {
		return err
	}

	return nil
}

func sendReplicate(task *models.Task, async bool) (*models.TaskResponse, error) {
	host, _ := beego.AppConfig.String("replicateUrl")
	url := host
	payload := new(bytes.Buffer)
	v := reflect.ValueOf(task.Input)
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}
	var input interface{}
	for _, key := range v.MapKeys() {
		value := v.MapIndex(key).Interface()
		keyStr := fmt.Sprintf("%v", key.Interface())
		logs.Debug("keyStr: value\n", keyStr, value)
		if keyStr == "input" {
			input = value
		}
	}
	//versions := map[string]string{
	//	"pulid":     "43d309c37ab4e62361e5e29b8e9e867fb2dcbcec77ae91206a8d95ac5dd451a0",
	//	"idm-vton":  "906425dbca90663ff5427624839572cc56ea7d380343d13e2a4c4b09d3f0c30f",
	//	"face-swap": "bc479d7d8ecc50eb83839af0c28210db75cac9c23837e2722028df4cddfafa22",
	//}
	parts := strings.Split(task.ApiPath, "/")
	lastElement := ""
	if len(parts) > 2 {
		lastElement = parts[len(parts)-1]
	}
	//version := versions[lastElement]
	//version, _ := beego.AppConfig.String(lastElement)
	//model, _ := config.GetSection(lastElement)
	model := supportModels[lastElement]
	taskReturn := &models.TaskReturn{
		Async: async,
	}
	taskResponse := &models.TaskResponse{
		Task: taskReturn,
	}
	if model == nil {
		task.Status = 3
		taskReturn.ExecError = "It`s not open yet."
		task.Error = taskResponse.Task
		_, err := mongo.Update(task)
		if err != nil {
			logs.Info("Update Task Error:", err)
		}
		return taskResponse, nil
	}
	data := models.ReplicateRequest{
		Version: model.Version,
		Input:   input,
		Stream:  model.Stream,
	}
	url = model.Url
	json.NewEncoder(payload).Encode(data)
	client := &http.Client{}
	request, err := http.NewRequest("POST", url, payload)
	if err != nil {
		setError(task, "sendReplicate request create error:"+err.Error())
		logs.Info("sendReplicate Error NewRequest request:", err)
		return nil, err
	}
	apikey, _ := beego.AppConfig.String("replicateToken")
	request.Header.Add("Authorization", "Bearer "+apikey)

	logs.Info("sendReplicate client sending request:")
	resp, err := client.Do(request)
	if err != nil {
		setError(task, "Task sending error:"+err.Error())
		logs.Info("sendReplicate Error sending request:", err)
		return nil, err
	}
	defer resp.Body.Close()
	logs.Info("sendReplicate resp code", resp.StatusCode)
	body, err := io.ReadAll(resp.Body)
	logs.Info("sendReplicate body", string(body))
	if resp.StatusCode == 200 && body == nil {
		setError(task, "Task response body null")
		logs.Info("sendReplicate Body reading response:", err)
		return nil, err
	}
	if err != nil {
		setError(task, "Task read response body error:"+err.Error())
		logs.Info("sendReplicate Error reading response:", err)
		return nil, err
	}
	var response models.ReplicateResponse
	if err = json.Unmarshal(body, &response); err != nil {
		setError(task, "Task response Unmarshal error:"+err.Error())
		logs.Info("sendReplicate Error Unmarshal response:", err)
		return nil, err
	}
	if response.Urls.Get != "" || (model.Stream && response.Urls.Stream != "") {
		//replicateTimeout, _ := beego.AppConfig.Int("replicateTimeout")
		//
		//timeout := time.After(time.Duration(replicateTimeout) * time.Minute)
		//
		//for {
		//	select {
		//	case <-timeout:
		//		logs.Info("Operation timed out")
		//		task.Status = 4 // 4表示超时状态
		//		taskReturn.TaskError = "Operation timed out"
		//		taskResponse.Task = taskReturn
		//		task.Error = taskResponse.Task
		//		_, err := mongo.Update(task)
		//		if err != nil {
		//			logs.Info("Update Task Error:", err)
		//		}
		//		return taskResponse, nil
		//	default:
		//		temp, err := getReplicate(response.Urls.Get)
		//		if err != nil {
		//			logs.Info("getReplicate Task Error:", err)
		//		}
		//		logs.Info("getReplicate Task temp:", temp)
		//		if temp != nil && temp.Status == "succeeded" {
		//			//todo 返回
		//			var output []string
		//			if str, ok := temp.Output.(string); ok {
		//				fmt.Println("i 是字符串类型，值为:", str)
		//				output = append(output, str)
		//			}
		//			if slice, ok := temp.Output.([]string); ok {
		//				fmt.Println("i 是字符串数组类型，值为:", slice)
		//				output = slice
		//			}
		//			if slice, ok := temp.Output.([]interface{}); ok {
		//				fmt.Println("i 是interface{}数组类型，值为:", slice)
		//				for _, value := range slice {
		//					if str, ok := value.(string); ok {
		//						fmt.Println("i 是字符串类型，值为:", str)
		//						output = append(output, str)
		//					}
		//				}
		//			}
		//			if slice, ok := temp.Output.(map[string]string); ok {
		//				fmt.Println("i 是map类型，值为:", slice)
		//				for _, value := range slice {
		//					output = append(output, value)
		//				}
		//			}
		//			task.Status = 2
		//			task.Output = output
		//			mongo.Update(task)
		//			taskResponse.Output = task.Output
		//			taskResponse.Task.IsSuccess = true
		//			taskResponse.Task.ExecCode = 200
		//
		//			return taskResponse, nil
		//		} else if temp != nil && temp.Error != nil {
		//			task.Status = 3
		//			if response.Error != nil {
		//				data, err := json.Marshal(response.Error)
		//				if err != nil {
		//					logs.Info("sendTask response.Task Unmarshal response:", err)
		//				}
		//				if data != nil && len(data) > 0 {
		//					taskReturn.ExecError = string(data)
		//				}
		//				taskResponse.Task = taskReturn
		//				task.Error = taskResponse.Task
		//			}
		//			_, err = mongo.Update(task)
		//			if err != nil {
		//				logs.Info("Update Task Error:", err)
		//			}
		//			return taskResponse, nil
		//		}
		//		time.Sleep(time.Second)
		//	}
		//}

		if async {
			go doGetReplicate(response.Urls.Get, task, taskResponse, taskReturn)
			return taskResponse, nil
		}
		doGetReplicate(response.Urls.Get, task, taskResponse, taskReturn)
		return taskResponse, nil
	} else {
		task.Status = 3
		if response.Error != nil {
			data, err := json.Marshal(response.Error)
			if err != nil {
				logs.Info("sendTask response.Task Unmarshal response:", err)
			}
			if data != nil && len(data) > 0 {
				taskReturn.ExecError = string(data)
			}
			taskResponse.Task = taskReturn
			task.Error = taskResponse.Task
		}
		_, err = mongo.Update(task)
		if err != nil {
			logs.Info("Update Task Error:", err)
		}

	}
	return taskResponse, nil
}

func doGetReplicate(url string, task *models.Task, taskResponse *models.TaskResponse, taskReturn *models.TaskReturn) {
	replicateTimeout, _ := beego.AppConfig.Int("replicateTimeout")

	timeout := time.After(time.Duration(replicateTimeout) * time.Minute)

	for {
		select {
		case <-timeout:
			logs.Info("Operation timed out")
			task.Status = 4 // 4表示超时状态
			taskReturn.TaskError = "Operation timed out"
			taskResponse.Task = taskReturn
			task.Error = taskResponse.Task
			_, err := mongo.Update(task)
			if err != nil {
				logs.Info("Update Task Error:", err)
			}
			return
		default:
			temp, err := getReplicate(url)
			if err != nil {
				logs.Info("getReplicate Task Error:", err)
			}
			logs.Info("getReplicate Task temp:", temp)
			if temp != nil && temp.Status == "succeeded" {
				//todo 返回
				var output []string
				if str, ok := temp.Output.(string); ok {
					fmt.Println("i 是字符串类型，值为:", str)
					output = append(output, str)
				}
				if slice, ok := temp.Output.([]string); ok {
					fmt.Println("i 是字符串数组类型，值为:", slice)
					output = slice
				}
				if slice, ok := temp.Output.([]interface{}); ok {
					fmt.Println("i 是interface{}数组类型，值为:", slice)
					for _, value := range slice {
						if str, ok := value.(string); ok {
							fmt.Println("i 是字符串类型，值为:", str)
							output = append(output, str)
						}
					}
				}
				if slice, ok := temp.Output.(map[string]string); ok {
					fmt.Println("i 是map类型，值为:", slice)
					for _, value := range slice {
						output = append(output, value)
					}
				}
				task.Status = 2
				task.Output = transferImages(output)
				mongo.Update(task)
				taskResponse.Output = task.Output
				taskResponse.Task.IsSuccess = true
				taskResponse.Task.ExecCode = 200
				return
			} else if temp != nil && temp.Error != nil {
				task.Status = 3
				if temp.Error != nil {
					data, err := json.Marshal(temp.Error)
					if err != nil {
						logs.Info("sendTask response.Task Unmarshal response:", err)
					}
					if data != nil && len(data) > 0 {
						taskReturn.ExecError = string(data)
					}
					taskResponse.Task = taskReturn
					task.Error = taskResponse.Task
				}
				_, err = mongo.Update(task)
				if err != nil {
					logs.Info("Update Task Error:", err)
				}
				return
			}
			time.Sleep(time.Second)
		}
	}

}

func getReplicate(url string) (*models.ReplicateResponse, error) {
	//host, _ := beego.AppConfig.String("replicateUrl")
	//url := host + task.ApiPath
	//payload := new(bytes.Buffer)
	//var input interface{}
	//if err := json.Unmarshal([]byte(task.Input), &input); err != nil {
	//	setError(task, "task.Input Unmarshal error:"+err.Error())
	//	logs.Info("sendTask task.Input Unmarshal response:", err)
	//	return nil, err
	//}
	//json.NewEncoder(payload).Encode(task.Input)
	client := &http.Client{}
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		logs.Info("getReplicate Error NewRequest request:", err)
		return nil, err
	}
	apikey, _ := beego.AppConfig.String("replicateToken")
	request.Header.Add("Authorization", "Bearer "+apikey)

	logs.Info("getReplicate client sending request:")
	resp, err := client.Do(request)
	if err != nil {
		logs.Info("getReplicate Error sending request:", err)
		return nil, err
	}
	defer resp.Body.Close()
	logs.Info("getReplicate resp code", resp.StatusCode)
	body, err := io.ReadAll(resp.Body)
	logs.Info("getReplicate body", string(body))
	if resp.StatusCode == 200 && body == nil {
		logs.Info("getReplicate Body reading response:", err)
		return nil, err
	}
	if err != nil {
		logs.Info("getReplicate Error reading response:", err)
		return nil, err
	}
	var response models.ReplicateResponse
	if err = json.Unmarshal(body, &response); err != nil {
		logs.Info("getReplicate Error Unmarshal response:", err)
		return nil, err
	}

	return &response, nil
}

func setError(task *models.Task, error string) {
	task.Status = 3
	task.Error = error
	_, err := mongo.Update(task)
	if err != nil {
		logs.Info("Update Task Error:", err)
	}
}

func (server *TaskController) Result() {
	taskId := server.GetString("excute_id")
	if taskId == "" {
		server.respond(models.NoRequestBody, "excute_id is null")
		return
	}
	task := models.Task{
		ExcuteId: taskId,
	}
	filter := map[string]string{
		"excute_id": taskId,
	}
	err := mongo.Read(&task, filter)
	if err != nil {
		logs.Info("List Error:", err)
		server.respond(models.BusinessFailed, err.Error())
	}
	server.respond(http.StatusOK, "", task)
}

func (server *TaskController) List() {
	body := server.Ctx.Input.RequestBody
	request := models.ListRequest{}
	err := json.Unmarshal(body, &request) //解析body中数据
	logs.Debug("Request", string(body), request)
	if err != nil {
		server.respond(models.NoRequestBody, err.Error())
		return
	}

	if request.Page == 0 {
		request.Page = 1
	}
	if request.Size == 0 {
		request.Size = 10
	}

	total, data, err := mongo.Query("Task", request.Page, request.Size, request.Filter)
	if err != nil {
		logs.Info("List Error:", err)
		server.respond(models.BusinessFailed, err.Error())
	}
	var tasks []models.Task
	for _, bsonD := range data {
		var task models.Task
		// 将 bson.D 转换为 bson.Raw
		bsonRaw, err := bson.Marshal(bsonD)
		if err != nil {
			server.respond(models.BusinessFailed, err.Error())
		}
		// 将 bson.Raw 解码为 User 结构体
		err = bson.Unmarshal(bsonRaw, &task)
		if err != nil {
			server.respond(models.BusinessFailed, err.Error())
		}
		if task.Input != nil {
			raw, err := bson.Marshal(task.Input)
			if err != nil {
				server.respond(models.BusinessFailed, err.Error())
			}
			// 将 bson.Raw 解码为 User 结构体
			var input map[string]interface{}
			err = bson.Unmarshal(raw, &input)
			if err != nil {
				server.respond(models.BusinessFailed, err.Error())
			}
			task.Input = input
		}

		if task.Error != nil {
			raw, err := bson.Marshal(task.Error)
			if err != nil {
				server.respond(models.BusinessFailed, err.Error())
			}
			// 将 bson.Raw 解码为 User 结构体
			var error1 models.TaskReturn
			err = bson.Unmarshal(raw, &error1)
			if err != nil {
				server.respond(models.BusinessFailed, err.Error())
			}
			task.Error = error1
		}

		tasks = append(tasks, task)
	}
	responseData := struct {
		Total int64       `json:"total"`
		Data  interface{} `json:"data,omitempty"`
	}{
		Total: total,
		Data:  tasks,
	}
	server.respond(http.StatusOK, "", responseData)
}
