package controller

import (
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"gorm.io/gorm"

	"github.com/tidbcloud/aitestsuit/configs"
	"github.com/tidbcloud/aitestsuit/consts"
	"github.com/tidbcloud/aitestsuit/dataset"
	"github.com/tidbcloud/aitestsuit/internal/dao"
	"github.com/tidbcloud/aitestsuit/internal/entity"
	"github.com/tidbcloud/aitestsuit/internal/tasks"
	"github.com/tidbcloud/aitestsuit/third_party/aipowersvc"
	"github.com/tidbcloud/aitestsuit/types"
)

func getLimitOffset(c *gin.Context) (int, int) {
	var pagination types.PaginationQuerySchema

	if err := c.ShouldBindQuery(&pagination); err != nil {
		pagination.Page = 1
	}

	return pagination.PerPage, (pagination.Page - 1) * pagination.PerPage
}

func Success(c *gin.Context, code int, message string, result interface{}) {
	c.JSON(code, gin.H{"message": message, "result": result})
}

func Fail(c *gin.Context, code int, message string, result interface{}) {
	c.JSON(code, gin.H{"message": message, "result": result})
}

func GetTestPlans(c *gin.Context) {
	ctx := c.Request.Context()
	limit, offset := getLimitOffset(c)
	tilog.Info(ctx, "limit and offset", zap.Int("limit", limit), zap.Int("offset", offset))

	total, plans, err := dao.GetAllTestPlansByPage(c, limit, offset)
	if err != nil {
		tilog.Error(ctx, "failed to get plans", zap.Error(err))
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	var result []gin.H
	for _, plan := range plans {
		casesCount := dao.MustCountCasesByPlanID(ctx, plan.ID)
		runsCount := dao.MustCountRunsByPlanID(ctx, plan.ID)
		reportsCount := dao.MustCountReportsByPlanID(ctx, plan.ID)
		var lastRunTime string
		lastRun, err := dao.GetLastTestReportByPlanID(ctx, plan.ID)
		if err == nil {
			lastRunTime = lastRun.CreatedAt.Format(consts.DATETIME_FMT)
		}
		aiModel, err := dao.GetAIModelByID(ctx, plan.AIModelID)
		if err != nil {
			tilog.Error(ctx, "failed to get ai model", zap.Error(err))
		}

		var (
			avgDuration      int64 // in ms
			totalDuration    float64
			validReportCount int
		)
		reports, err := dao.GetAllReportsByPlanID(ctx, plan.ID)
		if err == nil {
			for _, report := range reports {
				// ignore invalid status report
				if report.Status != consts.StatusOver {
					continue
				}

				// ignore report which complete too fast
				if report.EndAt.Sub(report.StartAt).Seconds() < 2 {
					continue
				}

				validReportCount++
				totalDuration += float64(report.EndAt.Sub(report.StartAt).Milliseconds())
			}
		}
		if validReportCount > 0 {
			avgDuration = int64(totalDuration / float64(validReportCount))
		}

		result = append(result, gin.H{
			"stats": gin.H{
				"cases_count":   casesCount,
				"runs_count":    runsCount,
				"reports_count": reportsCount,
				"avg_duration":  avgDuration,
			},
			"plan":     plan,
			"ai_model": aiModel,
			"last_run": lastRunTime,
		})
	}

	Success(c, http.StatusOK, "", gin.H{"plans": result, "total": total})
}

func GetTestCases(c *gin.Context) {
	limit, offset := getLimitOffset(c)
	ctx := c.Request.Context()

	var req types.GetCasesByPlanIDSchema
	if err := c.ShouldBindUri(&req); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	total, cases, err := dao.GetTestCasesByPlanIDAndPage(c, req.PlanID, limit, offset)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	result := []gin.H{}
	for _, testCase := range cases {
		var passRate float64
		caseResults, err := dao.GetAllTestCaseResultsByCaseID(ctx, testCase.ID)
		if err != nil {
			tilog.Error(ctx, "failed to get test case results by case id", zap.Error(err))
			continue
		} else {
			totalCount := float64(len(caseResults))
			var passCount float64
			for _, i := range caseResults {
				if i.ResultA == 1 {
					passCount += 1
				}
			}

			if totalCount > 0 {
				passRate = passCount / totalCount
			}
		}

		result = append(result, gin.H{
			"case":      testCase,
			"pass_rate": passRate,
		})
	}

	Success(c, http.StatusOK, "", gin.H{"cases": result, "total": total})
}

func GetTestCase(c *gin.Context) {
	var uriParams types.GetByIDSchema
	if err := c.ShouldBindUri(&uriParams); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	ctx := c.Request.Context()
	testCase, err := dao.GetTestCaseByID(ctx, uriParams.ID)
	if err != nil {
		Fail(c, http.StatusNotFound, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"case": testCase})
}

func DeleteTestCase(c *gin.Context) {
	var uriParams types.GetByIDSchema
	if err := c.ShouldBindUri(&uriParams); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	ctx := c.Request.Context()
	err := dao.DeleteTestCaseByID(ctx, uriParams.ID)
	if err != nil {
		Fail(c, http.StatusNotFound, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", nil)
}

func UpdateTestCase(c *gin.Context) {
	var payload types.UpdateTestCaseSchema
	if err := c.ShouldBindJSON(&payload); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}
	ctx := c.Request.Context()

	testCase, err := dao.GetTestCaseByID(ctx, payload.ID)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	testCase.PlanID = payload.PlanID
	testCase.Question = payload.Question
	testCase.Database = payload.Database
	testCase.Table = payload.Table
	testCase.SQL = payload.SQL
	testCase.ColumnNames = strings.Join(payload.ColumnNames, ",")
	testCase.CompareArgs = payload.CompareArgs

	err = dao.UpdateTestCase(ctx, &testCase)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"plan": testCase})
}

func GetPromptTemplates(c *gin.Context) {
	limit, offset := getLimitOffset(c)
	ctx := c.Request.Context()

	templates, err := dao.GetPromptTemplates(ctx, limit, offset)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"templates": templates})
}

func CreatePromptTemplate(c *gin.Context) {
	ctx := c.Request.Context()

	var req types.CreatePromptTemplateSchema
	if err := c.ShouldBindJSON(&req); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	template := entity.PromptTemplate{
		Version: req.Version,
		Content: req.Content,
	}

	if err := dao.CreatePromptTemplate(ctx, &template); err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"template": template})
}

func UpdatePromptTemplate(c *gin.Context) {
	ctx := c.Request.Context()
	var req types.GetByIDSchema
	if err := c.ShouldBindUri(&req); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	var templateReq types.CreatePromptTemplateSchema
	if err := c.ShouldBindJSON(&templateReq); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	template, err := dao.GetPromptTemplateByID(ctx, req.ID)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	template.Version = templateReq.Version
	template.Content = templateReq.Content
	if err = dao.UpdatePromptTemplate(ctx, &template); err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"template": template})
}

func ReleaseAIPowerPromptTemplate(c *gin.Context) {
	var req types.ReleaseAIModelPromptSchema
	if err := c.ShouldBindJSON(&req); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	ctx := c.Request.Context()
	template, err := dao.GetPromptTemplateByID(ctx, req.ID)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	_, err = aipowersvc.SavePromptTemplate(ctx, req.AIModel, req.Deployment, template)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", nil)
}

func GetAIPowerPromptTemplates(c *gin.Context) {
	filter := types.GetAllAIModelPromptsSchema{}
	if err := c.ShouldBindQuery(&filter); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}
	ctx := c.Request.Context()

	prompts, err := aipowersvc.GetAllPromptTemplate(ctx, filter.AIModel, filter.Deployment)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"prompts": prompts})
}

func GetCurrentAIPowerPromptTemplate(c *gin.Context) {
	filter := types.GetCurrentAIModelPromptSchema{}
	if err := c.ShouldBindQuery(&filter); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}
	ctx := c.Request.Context()

	prompt, err := aipowersvc.GetCurrentPromptTemplate(ctx, filter.AIModel, filter.Deployment)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"prompt": prompt})
}

func GetTestReports(c *gin.Context) {
	limit, offset := getLimitOffset(c)
	ctx := c.Request.Context()

	filter := types.GetReportsFilterSchema{}
	if err := c.ShouldBindQuery(&filter); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	total, reports, err := dao.GetTestReports(ctx, filter, limit, offset)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	result := []gin.H{}
	for _, report := range reports {
		var (
			generatingPassRate    float64
			executingPassRate     float64
			chartTypeMatchRate    float64
			chartTypeNonTableRate float64
		)

		plan, err := dao.GetTestPlanByID(ctx, report.PlanID)
		if err != nil {
			Fail(c, http.StatusInternalServerError, err.Error(), nil)
			return
		}
		run, err := dao.GetTestRunByID(ctx, report.TestRunID)
		if err != nil {
			Fail(c, http.StatusInternalServerError, err.Error(), nil)
			return
		}
		promptTemplate, err := dao.GetPromptTemplateByID(ctx, report.PromptTemplateID)
		if err != nil {
			Fail(c, http.StatusInternalServerError, err.Error(), nil)
			return
		}

		if report.Stats.TotalCount > 0 {
			generatingPassRate = float64(report.Stats.GeneratingPassCount) / float64(report.Stats.TotalCount)
			executingPassRate = float64(report.Stats.ExecutingPassCount) / float64(report.Stats.TotalCount)
			chartTypeMatchRate = float64(report.Stats.ChartTypeMatchCount) / float64(report.Stats.TotalCount)
			chartTypeNonTableRate = float64(report.Stats.NonTableChartCount) / float64(report.Stats.TotalCount)
		}

		result = append(result, gin.H{
			"generating_pass_rate":  generatingPassRate,
			"executing_pass_rate":   executingPassRate,
			"chart_type_match_rate": chartTypeMatchRate,
			"non_table_chart_rate":  chartTypeNonTableRate,
			"prompt_template":       promptTemplate,
			"report":                report,
			"plan":                  plan,
			"run":                   run,
			"start":                 report.StartAt.Format(consts.DATETIME_FMT),
		})
	}

	Success(c, http.StatusOK, "", gin.H{"reports": result, "total": total})
}

func GetTestReport(c *gin.Context) {
	limit, offset := getLimitOffset(c)
	ctx := c.Request.Context()
	var req types.GetByIDSchema
	if err := c.ShouldBindUri(&req); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}
	var filter types.GetReportDetailFilterSchema
	if err := c.ShouldBindQuery(&filter); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	testReport, err := dao.GetTestReport(ctx, req.ID)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	total, testCaseResults, err := dao.GetTestResultsByReportIDAndPage(ctx, req.ID, filter, limit, offset)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	var result []gin.H
	for _, caseResult := range testCaseResults {
		testCase, err := dao.GetTestCaseByID(ctx, caseResult.TestCaseID)
		if err != nil {
			Fail(c, http.StatusInternalServerError, err.Error(), nil)
			return
		}

		result = append(result, gin.H{
			"case_result": caseResult,
			"case":        testCase,
		})
	}

	Success(c, http.StatusOK, "", gin.H{
		"results":  result,
		"duration": testReport.EndAt.Sub(testReport.StartAt).Seconds(),
		"report":   testReport,
		"total":    total,
	})
}

func GetTestRuns(c *gin.Context) {
	limit, offset := getLimitOffset(c)
	ctx := c.Request.Context()

	runs, err := dao.GetTestRuns(ctx, limit, offset)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	var result []gin.H
	for _, run := range runs {
		result = append(result, gin.H{
			"run": run,
		})
	}

	Success(c, http.StatusOK, "", gin.H{"runs": result})
}

func GetTestRun(c *gin.Context) {

}

func GetTestStats(c *gin.Context) {
	limit, offset := getLimitOffset(c)
	ctx := c.Request.Context()

	var req types.GetStatsSchema
	if err := c.ShouldBindQuery(&req); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	runs, err := dao.GetTestRuns(ctx, limit, offset)
	if err != nil && err != gorm.ErrRecordNotFound {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	var result []gin.H
	planIDStrs := strings.Split(req.PlanIDs, ",")
	planIDs := []uint{}
	for _, i := range planIDStrs {
		id, _ := strconv.ParseUint(i, 10, 64)
		planIDs = append(planIDs, uint(id))
	}
	for _, run := range runs {
		reports := []entity.TestReport{}

		for _, planID := range planIDs {
			plan, err := dao.GetTestPlanByID(ctx, planID)
			if err != nil && err != gorm.ErrRecordNotFound {
				Fail(c, http.StatusInternalServerError, err.Error(), nil)
				return
			}

			report, err := dao.GetLastTestReportByRunIDAndPlanIDAndStatus(ctx, run.ID, plan.ID, consts.StatusOver)
			if err != nil && err != gorm.ErrRecordNotFound {
				Fail(c, http.StatusInternalServerError, err.Error(), nil)
				return
			}

			reports = append(reports, report)
		}

		result = append(result, gin.H{
			"run":     run,
			"reoprts": reports,
		})
	}

	Success(c, http.StatusOK, "", gin.H{"runs": result})
}

func CreateTestRun(c *gin.Context) {
	ctx := c.Request.Context()

	var req types.CreateTestRunSchema
	if err := c.ShouldBindJSON(&req); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	testRun := entity.TestRun{
		Name:   req.Name,
		Params: req.Params,
	}

	if err := dao.CreateTestRun(ctx, &testRun); err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", nil)
}

func CreateTestTask(c *gin.Context) {
	ctx := c.Request.Context()
	task := tasks.Task{}
	if err := c.ShouldBindJSON(&task); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	testRun, err := dao.GetTestRun(ctx, task.TestRunID)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	plan, err := dao.GetTestPlanByID(ctx, task.TestPlanID)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}
	cases, err := dao.GetAllTestCasesByPlanID(ctx, plan.ID)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}
	// prepare the report
	report := entity.TestReport{
		StartAt:          time.Now(),
		EndAt:            time.Now(),
		Status:           consts.StatusINIT,
		TokenUsed:        0,
		PlanID:           plan.ID,
		TestRunID:        testRun.ID,
		PromptTemplateID: task.PromptTemplateID,
		PassRate:         0,
	}
	err = dao.CreateTestReport(ctx, &report)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}
	var results []entity.TestCaseResult
	for _, testCase := range cases {
		result := entity.TestCaseResult{
			TestReportID:   report.ID,
			TestCaseID:     testCase.ID,
			Status:         consts.StatusINIT,
			Input:          "",
			Output:         "",
			ResultA:        0,
			ResultH:        0,
			ReturnSum:      0,
			ReturnColCount: 0,
			SQLResp:        entity.SQLResp{},
		}
		err = dao.CreateTestCaseResult(ctx, &result)
		if err != nil {
			Fail(c, http.StatusInternalServerError, err.Error(), nil)
			return
		}

		results = append(results, result)
	}
	cfg := configs.GetConfig()

	task.Params = testRun.Params
	task.ProjectID = int64(cfg.ProjectID)
	task.ClusterID = int64(cfg.ClusterID)
	task.Report = report
	task.Results = results

	tasks.EnqueueTask(ctx, task)
	Success(c, http.StatusOK, "", nil)
}

func CreateDataSet(c *gin.Context) {
	ctx := c.Request.Context()
	req := types.CreateDataSetSchema{}
	if err := c.ShouldBindJSON(&req); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}

	var err error
	if req.Name == consts.DataSetQuestion520 {
		err = dataset.GenQuestion(ctx, consts.DataSetQuestion520)
	} else if req.Name == consts.DataSetSpider {
		err = dataset.GenQuestion(ctx, consts.DataSetSpider)
	}

	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
	} else {
		Success(c, http.StatusOK, "", nil)
	}
}

func CreateTestPlan(c *gin.Context) {
	var payload types.CreateTestPlanSchema
	if err := c.ShouldBindJSON(&payload); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}
	ctx := c.Request.Context()

	testPlan := entity.TestPlan{
		Type:        payload.Type,
		Creator:     payload.Creator,
		Name:        payload.Name,
		Description: payload.Description,
		AIModelID:   payload.AIModelID,
	}

	err := dao.CreateTestPlan(ctx, &testPlan)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"plan": testPlan})
}

func CreateTestCase(c *gin.Context) {
	var payload types.CreateTestCaseSchema
	if err := c.ShouldBindJSON(&payload); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}
	ctx := c.Request.Context()

	testCase := entity.TestCase{
		PlanID:      payload.PlanID,
		Question:    payload.Question,
		Database:    payload.Database,
		Table:       payload.Table,
		SQL:         payload.SQL,
		ColumnNames: strings.Join(payload.ColumnNames, ","),
		CompareArgs: payload.CompareArgs,
	}

	err := dao.CreateTestCase(ctx, &testCase)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"plan": testCase})
}

func CreateAIModel(c *gin.Context) {
	var payload types.CreateAIModelSchema
	if err := c.ShouldBindJSON(&payload); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}
	ctx := c.Request.Context()

	aiModel := entity.AIModel{
		Name:        payload.Name,
		Description: payload.Description,
		ModelType:   payload.ModelType,
	}

	err := dao.CreateAIModel(ctx, &aiModel)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"ai_model": aiModel})
}

func UpdateAIModel(c *gin.Context) {
	var payload types.UpdateAIModelSchema
	if err := c.ShouldBindJSON(&payload); err != nil {
		Fail(c, http.StatusBadRequest, err.Error(), nil)
		return
	}
	ctx := c.Request.Context()

	aiModel, err := dao.GetAIModelByID(ctx, payload.ID)
	if err != nil {
		Fail(c, http.StatusNotFound, err.Error(), nil)
		return
	}

	aiModel.Name = payload.Name
	aiModel.Description = payload.Description
	aiModel.ModelType = payload.ModelType

	err = dao.UpdateAIModel(ctx, &aiModel)
	if err != nil {
		Fail(c, http.StatusInternalServerError, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"ai_model": aiModel})
}

func GetAIModels(c *gin.Context) {
	ctx := c.Request.Context()
	limit, offset := getLimitOffset(c)

	models, err := dao.GetAIModels(ctx, limit, offset)
	if err != nil {
		Fail(c, http.StatusNotFound, err.Error(), nil)
		return
	}

	Success(c, http.StatusOK, "", gin.H{"ai_models": models})
}
