package processor

import (
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/tidbcloud/observability/data-collector/pkg/common"
)

const MaxSize = 64 * 1024 * 1024 // 64MB

// ProcessCleanedData reads cleaned file data, aggregates it, writes to new files, and saves to disk
func (p *DataProcessor) ProcessCleanedData(readDir string, writeDir string, dataSource string, version string) error {
	var records = make(map[string]map[string]string)
	if dataSource == common.Chat2Query {
		pattern := fmt.Sprintf("%s/%s/%s/*.csv", readDir, dataSource, version)
		err := readFileToMapChat2Query(pattern, records)
		if err != nil {
			log.Printf("read file to map failed, err: %s", err.Error())
			return err
		}
		defer func() {
			err = DeleteFilesFunc(pattern)
			if err != nil {
				log.Printf("delete cleaned files failed, error: %v", err)
			}
		}()
	} else if dataSource == common.OSSInsight {
		patternA := fmt.Sprintf("%s/%s/%s/*_explorer_questions.csv", readDir, dataSource, version)
		patternB := fmt.Sprintf("%s/%s/%s/*_explorer_question_feedbacks.csv", readDir, dataSource, version)
		err := readFileToMapOssInsight(patternA, patternB, records)
		if err != nil {
			log.Printf("read file to map failed, err: %s", err.Error())
			return err
		}
		defer func() {
			err = DeleteFilesFunc(patternA)
			if err != nil {
				log.Printf("delete cleaned files %s failed, error: %v", patternA, err)
			}
			err = DeleteFilesFunc(patternB)
			if err != nil {
				log.Printf("delete cleaned files %s failed, error: %v", patternB, err)
			}
		}()
	}

	err := writeMapToCsvFile(writeDir, dataSource, version, records)
	if err != nil {
		log.Printf("write map to file failed, err: %s", err.Error())
	}
	return err
}

// MergeProcessedData merges all files generated by the processor in a day to a big file
func (p *DataProcessor) MergeProcessedData(bucket string, writeDir string, dataSource string, version string) error {
	now := time.Now()
	// TODO: read checkpoint
	// Process yesterday's data today to ensure no data is missing
	yesTime := now.AddDate(0, 0, -1)
	yesDate := yesTime.Format("20060102")

	// clusterIDMap[yesDate] records the cluster_ids processed yesterday through the store program
	var clusterIDMap map[string]map[string]bool
	if dataSource == common.Chat2Query {
		clusterIDMap = common.Chat2QueryClusterIDMap
	} else if dataSource == common.OSSInsight {
		clusterIDMap = map[string]map[string]bool{
			yesDate: {
				"default": true,
			},
		}
	}

	for clusterID, v := range clusterIDMap[yesDate] {
		if v {
			prefix := fmt.Sprintf("processed-data/%s/%s/%s/%s-", dataSource, version, clusterID, yesDate)
			objects, err := p.cli.DownloadFiles(bucket, prefix)
			if err != nil {
				log.Printf("get s3 objects failed, err: %s", err.Error())
				return err
			}
			log.Printf("start merge")
			for _, o := range objects {
				// Read each file content and append to new file
				path := fmt.Sprintf("%s/%s/%s/", writeDir, dataSource, version)
				err = writeObject(o, path, clusterID, yesDate)
				if err != nil {
					log.Printf("write s3 object to file within max size failed, err: %s", err.Error())
					return err
				}
			}
		}
		if _, ok := clusterIDMap[yesDate][clusterID]; ok {
			clusterIDMap[yesDate][clusterID] = false // Set to false, indicating the cluster_id has been processed
		}
	}
	delete(clusterIDMap, yesDate) // All data for yesDate has been processed
	return nil
}

// readFileToMap read fields of the row in the file into the map
// selectedID is an id selected to connect two tables, such as session_id or question_id
func readFileToMap(path string, mp map[string]map[string]string, selectedID string) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	r := csv.NewReader(file)
	// Read header
	head, err := r.Read()
	if err != nil {
		return err
	}
	var idIndex, querySqlIndex int
	querySqlExist := false
	for i, h := range head {
		if h == selectedID { // Find the selectedID index
			idIndex = i
		}
		if h == "query_sql" { // Find the query_sql index
			querySqlIndex = i
			querySqlExist = true
		}
	}
	for {
		row, err := r.Read()
		if err != nil && err != io.EOF {
			return err
		}
		if err == io.EOF {
			break
		}
		// Skip records when query_sql is empty
		if querySqlExist {
			if len(row[querySqlIndex]) == 0 { // TODO
				continue
			}
		}
		id := row[idIndex] // Get id
		if _, ok := mp[id]; !ok {
			mp[id] = make(map[string]string)
		}
		for i, h := range head {
			if h == selectedID {
				continue
			}
			mp[id][h] = row[i]
		}
	}
	return nil
}

var readFileToMapFunc = func(path string, mp map[string]map[string]string, selectedID string) error {
	return readFileToMap(path, mp, selectedID)
}

func readFileToMapChat2Query(pattern string, mp map[string]map[string]string) error {
	files, _ := filepath.Glob(pattern)
	log.Printf("read file %s, len: %v", pattern, len(files))
	// Merge the contents of the two tables into one map according to the session_id
	for _, file := range files {
		if err := readFileToMapFunc(file, mp, "session_id"); err != nil {
			return err
		}
	}
	return nil
}

func readFileToMapOssInsight(patternA string, patternB string, mp map[string]map[string]string) error {
	mpA := make(map[string]map[string]string)
	mpB := make(map[string]map[string]string)
	// Read explorer_questions files

	files, _ := filepath.Glob(patternA)
	log.Printf("read file %s, len: %v", patternA, len(files))
	for _, file := range files {
		err := readFileToMapFunc(file, mpA, "id")
		if err != nil {
			return err
		}
	}
	// Read explorer_question_feedbacks files
	files, _ = filepath.Glob(patternB)
	log.Printf("read file %s, len: %v", patternB, len(files))
	for _, file := range files {
		err := readFileToMapFunc(file, mpB, "question_id")
		if err != nil {
			return err
		}
	}

	// Merge two maps according to id==question_id
	for k, v := range mpA {
		mp[k] = make(map[string]string)
		for k1, v1 := range v {
			mp[k][k1] = v1
		}
		for k2, v2 := range mpB[k] {
			if k2 == "user_id" {
				continue
			}
			mp[k][k2] = v2
		}
		if p, ok := mp[k]["prompt"]; !ok || len(p) == 0 {
			mp[k]["prompt"] = common.DefaultOssInsightPrompt
		}
	}
	return nil
}

// writeMapToCsvFile selects the required fields content from the map according to the corresponding title,
// and then writes them to the file
func writeMapToCsvFile(dir string, dataSource string, version string, mp map[string]map[string]string) (err error) {
	title := common.GetDataTitle(dataSource)
	for k, v := range mp {
		// v is the content corresponding to a certain id
		row := make([]string, len(title))
		// row[0] saves the id such as session_id or question_id
		row[0] = k
		for i := 1; i < len(title); i++ {
			if a, ok := v[title[i]]; ok {
				row[i] = a
			}
		}

		outPath := fmt.Sprintf("%s/%s/%s/", dir, dataSource, version)
		clusterID, ok := v["cluster_id"]
		if !ok || len(clusterID) == 0 {
			clusterID = "default"
		}

		fn := fmt.Sprintf("%s/%s/%s/%s.csv", dir, dataSource, version, clusterID)
		log.Printf("write to file %s", fn)

		var w *csv.Writer
		var file *os.File
		defer file.Close()
		if _, err = os.Stat(fn); os.IsNotExist(err) {
			err = os.MkdirAll(outPath, 0700)
			if err != nil {
				return
			}
			file, err = os.Create(fn)
			if err != nil {
				return
			}
			w = csv.NewWriter(file)
			err = WriteCsvFileFunc(w, title)
			if err != nil {
				return
			}
		} else {
			file, err = os.OpenFile(fn, os.O_WRONLY|os.O_APPEND|os.O_CREATE, os.ModePerm)
			if err != nil {
				return
			}
			w = csv.NewWriter(file)
		}

		err = WriteCsvFileFunc(w, row)
		if err != nil {
			return
		}
		log.Printf("write to file %s done", fn)
	}
	return
}

func writeObject(object *s3.GetObjectOutput, path string, clusterID string, yesDate string) (err error) {
	r := csv.NewReader(object.Body)
	title, err := r.Read() // Read title
	if err != nil {
		return
	}
	for {
		row, err := r.Read()
		if err != nil && err != io.EOF {
			return err
		}
		if err == io.EOF {
			break
		}
		err = writeFileWithinMaxSize(path, clusterID, row, 0, yesDate, title, MaxSize)
		if err != nil {
			return err
		}
	}
	return
}

func getByte(str []string) int {
	num := 0
	for _, s := range str {
		num += len(s)
	}
	return num
}

func writeFileWithinMaxSize(path string, clusterID string, row []string, i int, yesDate string, title []string, maxSize int64) error {
	fn := fmt.Sprintf("%s%s-%s-%d.csv", path, clusterID, yesDate, i)
	var w *csv.Writer
	var file *os.File
	if _, err := os.Stat(fn); os.IsNotExist(err) {
		err = os.MkdirAll(path, 0700)
		if err != nil {
			return err
		}
		file, err = os.Create(fn)
		if err != nil {
			return err
		}
		w = csv.NewWriter(file)
		err = WriteCsvFileFunc(w, title)
		if err != nil {
			return err
		}
	} else {
		file, err = os.OpenFile(fn, os.O_WRONLY|os.O_APPEND|os.O_CREATE, os.ModePerm)
		if err != nil {
			return err
		}
		w = csv.NewWriter(file)
	}
	defer file.Close()

	fileInfo, _ := file.Stat()

	// Calculate the byte count of the row
	b := getByte(row)
	if fileInfo.Size()+int64(b)+4 <= maxSize {
		if err := WriteCsvFileFunc(w, row); err != nil {
			return err
		}
	} else {
		if err := writeFileWithinMaxSize(path, clusterID, row, i+1, yesDate, title, maxSize); err != nil {
			return err
		}
	}

	return nil
}