package workflow

import (
	"agent/src/common"
	"agent/src/message"
	"agent/src/utils/log"
	"bytes"
	"context"
	"strings"
	"sync"
	"time"
)

type handler func() error

type Task struct {
	ctx           context.Context
	client        common.WsClient
	request       any
	handler       handler
	task_id       uint64
	request_id    uint64
	pool          *TaskPool
	responseId    uint64
	response      chan any
	response_lock sync.Mutex
  	createTime    time.Time
}

func newTask(ctx context.Context, client common.WsClient, request_id uint64) *Task {
	return &Task{
		ctx       : ctx,
		client    : client,
		task_id   : common.NewTaskId(),
		request_id: request_id,
		response  : make(chan any, 1),
		createTime: time.Now(),
	}
}

func (task *Task) SendExecFail(errMsg string) {
	task.client.WriteJson(&message.ResponseExecFail{
		RequestId: task.request_id,
		Error    : errMsg,
	}, false)
}

func (task *Task) do() {
	defer func() {
		<-task.pool.sem
		task.pool.remove(task.client)
		if err := recover(); err != nil {
			log.Error("[Task] do panic", "err", err)
		}
	} ()

	if err := task.handler(); err != nil {
		log.Error("[Task] do", "err", err)
	}
}


func Handle(client common.WsClient, msg []byte) error {
	log.Info("[Task] handle", "msg", string(msg))
	
	var requestId uint64
	var errMsg string
	defer func() {
		if len(errMsg) > 0 {
			client.WriteJson(&message.ResponseExecFail{
				RequestId: requestId,
				Error: errMsg,
			}, false)
		}
	}()

	msg = bytes.TrimSpace(msg)
	cmd, err := message.ParseCmd(msg)
	if err != nil {
		errMsg = "invalid request format"
		return err
	}
	requestId = cmd.RequestId

	var execReq *message.RequestExec
	if strings.HasPrefix(cmd.Cmd, "exec_") {
		execReq, err = message.Parse[message.RequestExec](msg)
		if err != nil {
			errMsg = "invalid request format"
			return err
		}
		if err = execReq.Check(); err != nil {
			errMsg = err.Error()
			return err
		}
	}

	if cmd.Cmd == "ide_service" {
		req, err := message.Parse[message.MsgIDEService](msg)
		if err != nil {
			errMsg = "invalid ide request format"
			return err
		}
		err = client.NewWantResponse(req.RequestId, msg)
		if err != nil {
			errMsg = err.Error()
		}
		return err
	}

	task := newTask(client.GetCtx(), client, cmd.RequestId)
	var req any
	switch cmd.Cmd {
		case "list_model":
			return task.listModel(cmd.RequestId)
		case "exec_explain":
			req, err = message.Parse[message.RequestExecExplain](msg)
			task.request = req
			task.handler = task.explain
		case "exec_docstring":
			req, err = message.Parse[message.RequestExecDocstring](msg)
			task.request = req
			task.handler = task.docstring
		case "exec_fix":
			req, err = message.Parse[message.RequestExecFix](msg)
			task.request = req
			task.handler = task.fix
		case "exec_optimize":
			req, err = message.Parse[message.RequestExecFix](msg)
			task.request = req
			task.handler = task.optimize
		case "exec_unittest_recommend":
			req, err = message.Parse[message.RequestExecUnitTestRecommend](msg)
			task.request = req
			task.handler = task.unittest_recommend
		case "exec_chat":
			req, err = message.Parse[message.RequestExecChat](msg)
			task.request = req
			task.handler = task.exec_chat
		default:
			errMsg = "unknown cmd"
			return err
	}

	if err != nil {
		errMsg = err.Error()
		return err
	}

	if execReq != nil {
		pool := TaskPools[execReq.Model]
		err = pool.add(task)
		if err != nil {
			errMsg = err.Error()
			return err
		}
	}

	return nil
}
