package tasks

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync/atomic"
	"time"

	"github.com/gammazero/workerpool"
	aipower "github.com/tidbcloud/aipower/idl/pbgen/proto/api"
	"github.com/tidbcloud/dataservices/idl/pbgen/proto/dataapi"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"

	"github.com/tidbcloud/aitestsuit/configs"
	"github.com/tidbcloud/aitestsuit/consts"
	"github.com/tidbcloud/aitestsuit/internal/dao"
	"github.com/tidbcloud/aitestsuit/internal/entity"
	"github.com/tidbcloud/aitestsuit/third_party/aipowersvc"
	"github.com/tidbcloud/aitestsuit/third_party/dataapisvc"
)

var (
	msgCh = make(chan Task, 10240)
)

type Task struct {
	TestPlanID       uint          `json:"test_plan_id" binding:"required"`
	TestRunID        uint          `json:"test_run_id" binding:"required"`
	PromptTemplateID uint          `json:"prompt_template_id" binding:"required"`
	Params           entity.Params `json:"params"`
	ProjectID        int64         `json:"project_id"`
	ClusterID        int64         `json:"cluster_id"`

	Report  entity.TestReport       `json:"-"`
	Results []entity.TestCaseResult `json:"-"`
}

type Chart struct {
	ChartName string                 `json:"chartName"`
	Title     string                 `json:"title"`
	Options   map[string]interface{} `json:"options"`
}

type AIPowerCompletionResp struct {
	SQL   string `json:"sql"`
	Chart Chart  `json:"chart"`
}

func getSQLFromAIPower(ctx context.Context, completion string) string {
	completion = strings.ReplaceAll(completion, "\n", " ")

	var resp AIPowerCompletionResp
	if err := json.Unmarshal([]byte(completion), &resp); err != nil {
		tilog.Error(ctx, "failed to parse AIPower completion", zap.Error(err), zap.String("completion", completion))
		return completion
	}

	return resp.SQL
}

func getChartFromAIPower(ctx context.Context, completion string) Chart {
	completion = strings.ReplaceAll(completion, "\n", " ")

	var resp AIPowerCompletionResp
	if err := json.Unmarshal([]byte(completion), &resp); err != nil {
		tilog.Error(ctx, "failed to parse AIPower completion", zap.Error(err), zap.String("completion", completion))
	}

	return resp.Chart
}

func tryRunTask(ctx context.Context, model string, t Task, result entity.TestCaseResult) (*aipower.ChatBotRsp, *dataapi.ExecSqlResp, error) {
	testCase, err := dao.GetTestCaseByID(ctx, result.TestCaseID)
	if err != nil {
		tilog.Error(ctx, "failed to get test case", zap.Error(err), zap.Uint("caseID", result.TestCaseID))
		return nil, nil, err
	}

	template, err := dao.GetPromptTemplateByID(ctx, t.PromptTemplateID)
	if err != nil {
		tilog.Error(ctx, "failed to load prompt template", zap.Error(err), zap.Uint("templateID", t.PromptTemplateID))
		return nil, nil, err
	}
	tilog.Info(ctx, "using prompt template", zap.Any("template", template))

	cfg := configs.GetConfig()
	aiPowerResp, err := aipowersvc.GenSQLByTestCase(ctx, cfg.ProjectID, cfg.ClusterID, model, testCase, template)
	if err != nil {
		tilog.Error(ctx, "failed to get response from chat2data", zap.Error(err), zap.Any("task", t), zap.Any("aiPowerResp", aiPowerResp))
		return aiPowerResp, nil, err
	}
	tilog.Info(ctx, "get reply from aipower", zap.Any("aiPowerResp", aiPowerResp))
	if aiPowerResp.Code != 0 {
		return aiPowerResp, nil, errors.New(aiPowerResp.BaseResp.ErrMsg)
	}

	sql := getSQLFromAIPower(ctx, aiPowerResp.Completion)

	sqlResp, err := dataapisvc.TryExecSQL(ctx, cfg.ProjectID, cfg.ClusterID, testCase.Database, sql)
	if err != nil {
		return aiPowerResp, sqlResp, fmt.Errorf("[exec SQL from AI]: %w", err)
	}
	tilog.Info(ctx, "get reply from dataapi", zap.Any("sqlResp", sqlResp))

	return aiPowerResp, sqlResp, nil
}

func compareSQLResp(stdSQLResp *dataapi.ExecSqlResp, sqlResp *dataapi.ExecSqlResp) (bool, error) {
	if sqlResp.RowNum != stdSQLResp.RowNum {
		return false, fmt.Errorf(
			"SQL result does not match, should be %dR/%dC, got %dR/%dC",
			stdSQLResp.RowNum, len(stdSQLResp.Columns), sqlResp.RowNum, len(sqlResp.Columns),
		)
	}

	columnNames := map[string]bool{}
	for _, i := range sqlResp.Columns {
		columnNames[i.Name] = true
	}

	for _, i := range stdSQLResp.Columns {
		if !columnNames[i.Name] {
			return false, fmt.Errorf("SQL result does not match, should contain column %s but not exists", i.Name)
		}
	}

	return true, nil
}

func compareChartName(chartInfo Chart, testCase *entity.TestCase) (bool, error) {
	if chartInfo.ChartName == testCase.CompareArgs.Chart.ChartName {
		return true, nil
	} else {
		return false, fmt.Errorf("chart does not match, expect %+v, got %+v", testCase.CompareArgs.Chart, chartInfo)
	}
}

func compareTaskResult(
	ctx context.Context, testCaseID uint, model string,
	aiPowerResp *aipower.ChatBotRsp, sqlResp *dataapi.ExecSqlResp, metrics *ResultMetrics,
) (bool, error) {
	if sqlResp == nil {
		return false, nil
	}

	if aiPowerResp != nil {
		metrics.SQLGeneratedCount.Add(1)
	}
	if sqlResp != nil {
		metrics.SQLExecutedPassCount.Add(1)
	}

	testCase, err := dao.GetTestCaseByID(ctx, testCaseID)
	if err != nil {
		tilog.Error(ctx, "failed to get test case", zap.Error(err), zap.Uint("caseID", testCaseID))
		return false, err
	}

	tilog.Info(ctx, "check if sql in test case is not empty", zap.Any("sql", testCase.SQL))
	// both chat2sql, chat2chart needs check it if SQL provided
	if testCase.SQL != "" {
		tilog.Info(ctx, "execute SQL to compare result", zap.Any("testCase", testCase))
		cfg := configs.GetConfig()
		stdSQLResp, err := dataapisvc.TryExecSQL(ctx, cfg.ProjectID, cfg.ClusterID, testCase.Database, testCase.SQL)
		if err != nil {
			return false, fmt.Errorf("[exec SQL from std case]: %w", err)
		}

		pass, err := compareSQLResp(stdSQLResp, sqlResp)
		if err != nil || !pass {
			return pass, err
		}
	}

	tilog.Info(ctx, "check if chart name in compare args is not empty", zap.Any("compareArgs", testCase.CompareArgs))
	// only chat2chart needs check it
	if model == consts.ModelChat2Chart && testCase.CompareArgs.Chart.ChartName != "" {
		chart := getChartFromAIPower(ctx, aiPowerResp.Completion)

		tilog.Info(ctx, "check if chart name is table", zap.Any("chart", chart), zap.String("completion", aiPowerResp.Completion))
		if chart.ChartName != "Table" {
			metrics.ChartTypeNonTableCount.Add(1)
		}

		isMatch, err := compareChartName(chart, &testCase)
		if err != nil || !isMatch {
			return isMatch, err
		}

		tilog.Info(ctx, "chart name match", zap.Any("chart", chart))
		metrics.ChartTypeMatchCount.Add(1)
	}

	return true, nil
}

type ResultMetrics struct {
	SQLGeneratedCount       atomic.Int32
	SQLExecutedPassCount    atomic.Int32
	ResultComparedPassCount atomic.Int32
	ChartTypeMatchCount     atomic.Int32
	ChartTypeNonTableCount  atomic.Int32
	TotalTokenUsed          atomic.Int32
}

func runTask(ctx context.Context, t Task) {
	tilog.Info(ctx, "running task", zap.Any("task", t))

	report := t.Report
	report.Status = consts.StatusRunning
	if err := dao.UpdateTestReport(ctx, &report); err != nil {
		tilog.Error(ctx, "failed to update report status", zap.Error(err))
	}

	testPlan, err := dao.GetTestPlanByID(ctx, t.TestPlanID)
	if err != nil {
		tilog.Error(ctx, "failed to get test plan", zap.Error(err))
		return
	}

	aiModel, err := dao.GetAIModelByID(ctx, testPlan.AIModelID)
	if err != nil {
		tilog.Error(ctx, "failed to get ai model", zap.Error(err))
		return
	}

	model := aiModel.ModelType

	var metrics ResultMetrics
	wp := workerpool.New(10)

	for _, result := range t.Results {
		t := t
		result := result

		wp.Submit(func() {
			aiPowerResp, sqlResp, err1 := tryRunTask(ctx, model, t, result)
			if err1 != nil {
				result.ErrorLog = err1.Error()
			}

			// record the result
			pass, err2 := compareTaskResult(ctx, result.TestCaseID, model, aiPowerResp, sqlResp, &metrics)
			if err2 != nil {
				if len(result.ErrorLog) > 0 {
					result.ErrorLog += "\n"
				}
				result.ErrorLog += fmt.Sprintf("%s", err2)
			}
			if pass {
				metrics.ResultComparedPassCount.Add(1)
				result.ResultA = 1
			}

			result.Status = consts.StatusOver

			if aiPowerResp != nil {
				result.Duration = aiPowerResp.Duration
				result.Input = aiPowerResp.Prompt
				result.TokenUsed = int(aiPowerResp.TokenUsed)
				result.Output = getSQLFromAIPower(ctx, aiPowerResp.Completion)
			}

			if sqlResp != nil {
				result.ReturnSum = int64(len(sqlResp.Rows))
				result.ReturnColCount = int64(len(sqlResp.Columns))
				result.SQLResp = entity.SQLResp{DataAPIExecSQLResp: sqlResp}
			}

			metrics.TotalTokenUsed.Add(int32(result.TokenUsed))

			err1 = dao.UpdateTestCaseResult(ctx, &result)
			tilog.Info(ctx, "update test case result", zap.Error(err1), zap.Uint("case_result_id", result.ID), zap.Any("result", result))
		})
	}

	wp.StopWait()

	// update test report statistics
	tilog.Info(ctx, "statistics", zap.Any("metrics", &metrics))
	report.Stats = entity.Stats{
		GeneratingPassCount: int(metrics.SQLGeneratedCount.Load()),
		ExecutingPassCount:  int(metrics.SQLExecutedPassCount.Load()),
		ComparingPassCount:  int(metrics.ResultComparedPassCount.Load()),
		ChartTypeMatchCount: int(metrics.ChartTypeMatchCount.Load()),
		NonTableChartCount:  int(metrics.ChartTypeNonTableCount.Load()),
		TotalCount:          len(t.Results),
	}
	report.PassRate = float64(report.Stats.ComparingPassCount) / float64(report.Stats.TotalCount)
	report.Status = consts.StatusOver
	report.TokenUsed = int(metrics.TotalTokenUsed.Load())
	report.EndAt = time.Now()
	err = dao.UpdateTestReport(ctx, &report)
	if err != nil {
		tilog.Error(ctx, "failed to update test report", zap.Error(err))
	}
}

func EnqueueTask(ctx context.Context, t Task) {
	tilog.Info(ctx, "enqueue task", zap.Any("task", t))
	msgCh <- t
}

func StartWorker(ctx context.Context) {
	tilog.Info(ctx, "running worker")

	for {
		t := <-msgCh
		go runTask(ctx, t)
	}
}