package service

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/go-resty/resty/v2"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"

	"github.com/tidbcloud/aitestsuit/configs"
)

var restClient = resty.New()

type Chat2DataResp struct {
	Code           int    `json:"code"`
	Message        string `json:"message"`
	GenSQL         string `json:"gen_sql"`
	CallOpenaiTime string `json:"call_openai_time"`
	StartTime      string `json:"start_time"`
	EndTime        string `json:"end_time"`
	ExecuteTime    string `json:"execute_time"`
	Query          string `json:"query"`
	RowCount       int    `json:"row_count"`
	Columns        []struct {
		Col      string `json:"col"`
		DataType string `json:"data_type"`
		Nullable bool   `json:"nullable"`
	} `json:"columns"`
	Rows      [][]string `json:"rows"`
	Limit     int        `json:"limit"`
	ChartInfo struct {
		ChartName string `json:"chartName"`
		Title     string `json:"title"`
		Options   struct {
			X string   `json:"x"`
			Y []string `json:"y"`
		} `json:"options"`
	} `json:"chart_info"`
	CostTime   string `json:"cost_time"`
	TokenUsed  int    `json:"token_used"`
	FullPrompt string `json:"full_prompt"`
}

func Chat2Data(ctx context.Context, projectID, clusterID int64, database string, question string, promptTemplate, promptTemplateVersion string) (*Chat2DataResp, error) {
	config := configs.GetConfig()

	url := config.DataAPIURL
	payload := map[string]string{
		"project_id":              fmt.Sprintf("%d", projectID),
		"cluster_id":              fmt.Sprintf("%d", clusterID),
		"database":                database,
		"instruction":             question,
		"prompt_template":         promptTemplate,
		"prompt_template_version": promptTemplateVersion,
	}

	resp, err := restClient.R().SetBody(&payload).Post(url)
	if err != nil {
		tilog.Error(ctx, "failed to generate sql from dataapi", zap.Error(err), zap.Any("payload", payload))
		return nil, err
	}

	var result Chat2DataResp
	if err = json.Unmarshal([]byte(resp.String()), &result); err != nil {
		tilog.Error(ctx, "failed to unmarshal result from dataapi", zap.Error(err), zap.Any("resp", resp.String()))
		return nil, err
	}

	tilog.Info(ctx, "generate sql from dataapi", zap.Any("resp", resp.String()), zap.Any("payload", payload), zap.Any("result", result))

	if result.Code != http.StatusOK {
		return &result, fmt.Errorf("code %d, message %s", result.Code, result.Message)
	}

	return &result, nil
}