package processor

import (
	"bufio"
	"database/sql"
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go-v2/service/s3"
	_ "github.com/mattn/go-sqlite3"
)

// GenerateDB generates db based on the schema in the prompt
// bucket = gpt
// dir = /mnt/data/constructed-db
func (p *DataProcessor) GenerateDB(bucket string, dir string, dataSource string, version string) error {
	outPath := fmt.Sprintf("%s/%s/%s", dir, dataSource, version)
	fn := fmt.Sprintf("%s/checkpoint.json", outPath)
	cp, err := GetCheckPointFunc(fn) // If there is no checkpoint, it means that the operation has not started under this version, and all the files need to be pulled again, with the version as the prefix
	timeNow := time.Now()
	ts := timeNow.Unix()
	fmt.Println(err)
	if err != nil {
		if os.IsNotExist(err) {
			err = p.DownloadAndConstructDB(bucket, dataSource, version, outPath, "20")
			if err != nil {
				return err
			}
			SaveCheckPoint(fn, ts)
			return nil
		} else {
			return err
		}
	}
	// get start time point
	cpTs := cp.Timestamp
	for {
		if cpTs+86400 > ts {
			break
		}
		tm := time.Unix(cpTs, 0)
		date := tm.Format("20060102")
		err = p.DownloadAndConstructDB(bucket, dataSource, version, outPath, date)
		if err != nil {
			return err
		}
		cpTs += 86400
		SaveCheckPoint(fn, cpTs) // update checkpoint
	}
	return nil
}

func (p *DataProcessor) DownloadAndConstructDB(bucket string, dataSource string, version string, outPath string, dateTime string) error {
	prefix := fmt.Sprintf("dataset/joined-data/%s/%s/%s", dataSource, version, dateTime)
	objects, err := p.cli.DownloadFiles(bucket, prefix)
	if err != nil {
		log.Printf("download data failed, err: %s\n", err.Error())
		return err
	}
	for _, o := range objects {
		err = constructDB(o, outPath)
		if err != nil {
			log.Printf("clean data failed, err: %s\n", err.Error())
			return err
		}
	}
	return nil
}

func constructDB(object *s3.GetObjectOutput, outPath string) error {
	reader := csv.NewReader(object.Body)
	title, _ := reader.Read()
	index := getIndex(title)
	for {
		row, err := reader.Read()
		if err != nil && err != io.EOF {
			return err
		}
		if err == io.EOF {
			break
		}
		prompt := row[index["prompt"]]
		id := row[index["id"]] // todo id
		sql := row[index["query_sql"]]
		err = writeFile(fmt.Sprintf("%s/%s/querysql.txt", outPath, id), sql)
		if err != nil {
			return err
		}
		tables, err := UnmarshalPromptToDBFunc(prompt)
		if err != nil {
			return err
		}
		for _, t := range tables {
			err = createTable(t, fmt.Sprintf("%s/%s", outPath, id))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

type Table struct {
	DBName    string
	TableName string
	Columns   []string
	Type      []string
}

var UnmarshalPromptToDBFunc = func(prompt string) ([]Table, error) {
	return UnmarshalPromptToDB(prompt)
}

func UnmarshalPromptToDB(prompt string) ([]Table, error) {
	prefixIndex := strings.Index(prompt, "database schema")
	suffixIndex := strings.Index(prompt, ". The SQL statement you generate")
	if suffixIndex < 0 {
		suffixIndex = strings.Index(prompt, ". The question is:")
	}
	str := prompt[prefixIndex+len("database schema: ") : suffixIndex]
	tables := strings.Split(str, "Table ")
	result := make([]Table, 0)
	for _, j := range tables {
		if j == "" {
			continue
		}
		result = append(result, getTableInfo(j))
	}
	return result, nil
}

func getTableInfo(describe string) Table {
	describe = strings.ReplaceAll(describe, "`", "")
	indexA := strings.Index(describe, ".")
	indexB := strings.Index(describe, ",")
	indexC := strings.Index(describe, "columns = [")
	indexD := strings.Index(describe, "]")
	db := describe[:indexA]
	table := describe[indexA+1 : indexB]
	column := describe[indexC+len("columns = [") : indexD]
	columns := strings.Split(column, ", ")
	return Table{
		DBName:    db,
		TableName: table,
		Columns:   columns,
		Type:      getTypes(columns),
	}
}

func createTable(table Table, dir string) error {
	// create database
	if _, err := os.Stat(fmt.Sprintf("%s/%s.sqlite", dir, table.DBName)); os.IsNotExist(err) {
		err = os.MkdirAll(dir, 0700)
		if err != nil {
			return err
		}
	}
	db, err := sql.Open("sqlite3", fmt.Sprintf("%s/%s.sqlite", dir, table.DBName))
	fmt.Println("errr", err)
	if err != nil {
		return err
	}
	// create table
	sql := createTableSql(table.TableName, table.Columns, table.Type)
	re, err := db.Exec(sql)
	if re == nil {
		fmt.Println("nnnnnil")
	}
	return err
}

func createTableSql(table string, columns []string, types []string) string {
	col := fmt.Sprintf("\n %s %s PRIMARY KEY", columns[0], types[0])
	for i := 1; i < len(columns); i++ {
		col += fmt.Sprintf(",\n %s %s NULL", columns[i], types[i])
	}
	sql := fmt.Sprintf(sqlTemplate, table, col)
	return sql
}

var sqlTemplate = "CREATE TABLE IF NOT EXISTS %s(%s);"

func getTypes(columns []string) []string {
	types := make([]string, len(columns))
	for i, c := range columns {
		if strings.Contains(c, "age") {
			types[i] = "INT"
		} else {
			types[i] = "TEXT"
		}
	}
	return types
}

func getIndex(title []string) map[string]int {
	indexMap := make(map[string]int)
	for i, j := range title {
		indexMap[j] = i
	}
	return indexMap
}

func writeFile(fn string, querySql string) error {
	file, err := os.OpenFile(fn, os.O_WRONLY|os.O_CREATE, os.ModePerm)
	if err != nil {
		if err != nil {
			return err
		}
	}
	defer file.Close()
	write := bufio.NewWriter(file)
	_, err = write.WriteString(querySql)
	if err != nil {
		return err
	}
	write.Flush()
	return nil
}