package processor

import (
	"bytes"
	"fmt"
	"os"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	s3api "github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"
	"github.com/tidbcloud/observability/data-collector/pkg/common"
	mock_data_processor "github.com/tidbcloud/observability/data-collector/test/mockdataprocessor"
)

func TestUnmarshalPromptToDB(t *testing.T) {
	dir, _ := os.Getwd()
	println(dir)
	prompt := "Act as a MySQL expert, you are tasked with writing a SQL statement to answer question based on the following database schema: Table sample_data.sold_car_orders, columns = [`name`, `year`, `selling_price`, `km_driven`, `fuel`, `seller_type`, `transmission`, `owner`] Table sample_data.global_fortune_500_2018_2022, columns = [`rank`, `company_name`, `country`, `employees_num`, `previous_rank`, `revenues_million`, `revene_change_percentage`, `profit_million`, `profit_change_percentage`, `asset_million`, `year`] Table sample_data.github_events, columns = [`id`, `type`, `created_at`, `repo_id`, `repo_name`, `actor_id`, `actor_login`, `language`, `additions`, `deletions`, `action`, `number`, `commit_id`, `comment_id`, `org_login`, `org_id`, `state`, `closed_at`, `comments`, `pr_merged_at`, `pr_merged`, `pr_changed_files`, `pr_review_comments`, `pr_or_issue_id`, `event_day`, `event_month`, `event_year`, `push_size`, `push_distinct_size`, `creator_user_login`, `creator_user_id`, `pr_or_issue_created_at`] Table sample_data.imdb_movie_ratings, columns = [`rank`, `movie_name`, `year`, `certificate`, `runtime_in_min`, `genre`, `metascore`, `gross_in_$_m`, `rating_from_10`]. The question is: \" Find the most popular repo from GitHub events\". Let's think step by step and use the following answer template to generate an json file: { 1. Output the revised question, including all necessary keywords. Marked as RQ. {\"RQ\": string}. 2. List all necessary tables involved in answering the RQ. Mark as NT. {\"NT\": string[]} 3. List all necessary columns of NT involved in answering the RQ, do not introduce non-existent columns. Marked as NC. {\"NC\": {table: string, Columns: string[]}} 4. Using best practices of MySQL to write a single SQL statement to answer the RQ. {\"sql\": string} 5. Based on the following rules, rewrite the SQL statement. If no SQL statement can be generated, output the message 'No SQL Generated'. {\"rewriteSql\": string} ''' - Recheck the column names listed in NC, replace the corresponding column names in the SQL statement. - Only use the table.column format in the SQL statement when multiple tables mentioned in the NT list. - Ensure that columns are referenced properly and only exist in the appropriate table. - Enclose all table names and column names with backtick. ''' }\n"
	tables, err := UnmarshalPromptToDB(prompt)
	fmt.Println(err)
	fmt.Println("len", len(tables))

	for _, t := range tables {
		createTable(t, "./test/id")
	}
}

func TestDataProcessor_GenerateDB(t *testing.T) {
	ctl := gomock.NewController(t)
	s3 := mock_data_processor.NewMocks3Client(ctl)
	mockS3 := newMockS3Client(s3)
	processor := newMockDataProcessor(mockS3)
	t.Run("no checkpoint", func(t *testing.T) {
		s3.EXPECT().ListObjects(gomock.Any(), gomock.Any(), gomock.Any()).Return(&s3api.ListObjectsOutput{
			Contents: []types.Object{
				{
					Key: aws.String("20230308-12-0-timestamp_explorer_questions.csv"),
				},
			},
			Prefix: aws.String("test/data"),
		}, nil).Times(1)
		buf := make([]byte, 1)
		tmp := CBuffer{bytes.NewBuffer(buf)}
		s3.EXPECT().GetObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(&s3api.GetObjectOutput{
			Body:          &tmp,
			ContentLength: 1,
		}, nil).Times(1)

		err := processor.GenerateDB("test", "test", "ossinsight", "v1")
		defer deleteData("./test")
		assert.NoError(t, err)
	})

	t.Run("checkpoint", func(t *testing.T) {
		ts := time.Now().Unix() - 86500
		os.MkdirAll("test/ossinsight/v1/", 0700)
		SaveCheckPoint("test/ossinsight/v1/checkpoint.json", ts)
		s3.EXPECT().ListObjects(gomock.Any(), gomock.Any(), gomock.Any()).Return(&s3api.ListObjectsOutput{
			Contents: []types.Object{
				{
					Key: aws.String("20230308-12-0-timestamp_explorer_questions.csv"),
				},
			},
			Prefix: aws.String("test/data"),
		}, nil).Times(1)
		buf := make([]byte, 1)
		tmp := CBuffer{bytes.NewBuffer(buf)}
		s3.EXPECT().GetObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(&s3api.GetObjectOutput{
			Body:          &tmp,
			ContentLength: 1,
		}, nil).Times(1)

		err := processor.GenerateDB("test", "test", "ossinsight", "v1")
		defer deleteData("./test")
		assert.NoError(t, err)
	})

	t.Run("download err", func(t *testing.T) {
		s3.EXPECT().ListObjects(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, fmt.Errorf("error")).Times(1)
		err := processor.GenerateDB("test", "test", "ossinsight", "v1")
		assert.Error(t, err)
	})
	t.Run("get checkpoint err", func(t *testing.T) {
		GetCheckPointFunc = func(fn string) (*common.CheckPoint, error) {
			return nil, fmt.Errorf("error")
		}
		defer func() {
			GetCheckPointFunc = func(fn string) (*common.CheckPoint, error) {
				cp, err := GetCheckPoint(fn)
				return cp, err
			}
		}()

		err := processor.GenerateDB("test", "test", "ossinsight", "v1")
		assert.Error(t, err)
	})
}