package schema_utils

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"

	_ "github.com/mattn/go-sqlite3"
	"github.com/pingcap/log"
	"github.com/pkg/errors"
	"go.uber.org/zap"
)

type ColumnInfo struct {
	Name    string `json:"name"`
	Typ     string `json:"type"`
	Comment string `json:"comment"`
}

type TableInfo struct {
	DB      string                 `json:"db"`
	Table   string                 `json:"table"`
	Columns map[string]*ColumnInfo `json:"columns"`
	Comment string                 `json:"comment"`
	Tags    []string               `json:"-"`
}

func NewTableInfo(db, table, comment string) *TableInfo {
	return &TableInfo{
		DB:      db,
		Table:   table,
		Comment: comment,
		Columns: make(map[string]*ColumnInfo),
	}
}

func (t *TableInfo) AddColumn(name, typ, comment string) {
	if t.Columns == nil {
		t.Columns = make(map[string]*ColumnInfo)
	}
	if _, ok := t.Columns[name]; ok {
		log.Warn("duplicate column name", zap.String("table", t.Table), zap.String("column", name))
	}

	t.Columns[name] = &ColumnInfo{
		Name:    name,
		Typ:     typ,
		Comment: comment,
	}
}

func (t *TableInfo) CreateTableStatement(withComment bool) string {
	CreateTableStmt := ""
	stmt := fmt.Sprintf("Create table `%s` (", t.Table)

	names := make([]string, 0, len(t.Columns))
	for name := range t.Columns {
		names = append(names, name)
	}
	sort.Strings(names)

	for index, name := range names {
		colInfo := t.Columns[name]
		if index == 0 {
			stmt += fmt.Sprintf("`%s` %s", name, colInfo.Typ)
		} else {
			stmt += fmt.Sprintf(", `%s` %s", name, colInfo.Typ)
		}
		if withComment && len(colInfo.Comment) > 0 {
			comment := strings.Replace(colInfo.Comment, "'", "\\'", -1)
			stmt += fmt.Sprintf(" COMMENT '%s'", comment)
		}
	}

	if withComment && len(t.Comment) > 0 {
		CreateTableStmt = stmt + fmt.Sprintf(") COMMENT '%s';", t.Comment)
	} else {
		CreateTableStmt = stmt + ");"
	}

	return CreateTableStmt
}

type SpiderTableJson struct {
	DBID               string          `json:"db_id"`
	TableNamesOriginal []string        `json:"table_names_original"`
	TableNames         []string        `json:"table_names"`
	PrimayKeys         []int           `json:"primary_keys"`
	Columns            [][]interface{} `json:"column_names"`
	ColumnsOriginal    [][]interface{} `json:"column_names_original"`
	ColumnTypes        []string        `json:"column_types"`
	ForeignKeys        [][]int         `json:"foreign_keys"`
}

type DB interface {
	FetchTableInfo(context.Context, string, string, string) (*TableInfo, error)
}

type DBInfo struct {
	db     string
	tables map[string]*TableInfo
}

func (d *DBInfo) GenerateTableJson(w io.Writer) error {
	spiderTableJson := &SpiderTableJson{
		DBID:               d.db,
		TableNamesOriginal: make([]string, len(d.tables)),
		TableNames:         make([]string, len(d.tables)),
		ForeignKeys:        make([][]int, 0, 1),
		PrimayKeys:         make([]int, 0, 1),
		ColumnsOriginal:    make([][]interface{}, 0, 1),
		Columns:            make([][]interface{}, 0, 1),
	}

	index := 0
	for tbl, info := range d.tables {
		spiderTableJson.TableNames[index] = tbl
		spiderTableJson.TableNamesOriginal[index] = tbl

		for _, colInfo := range info.Columns {
			spiderTableJson.ColumnTypes = append(spiderTableJson.ColumnTypes, colInfo.Typ)
			spiderTableJson.Columns = append(spiderTableJson.Columns, []interface{}{index, colInfo.Name})
			spiderTableJson.ColumnsOriginal = append(spiderTableJson.ColumnsOriginal, []interface{}{index, colInfo.Name})
		}

		index++
	}

	tableJson, err := json.Marshal(spiderTableJson)
	if err != nil {
		return errors.Errorf("Fail to encode table json of %s: %v", d.db, err)
	}

	n, err := w.Write(tableJson)
	if err != nil {
		return errors.Errorf("Fail to write table json of %s: %v", d.db, err)
	}
	if n != len(tableJson) {
		return errors.Errorf("Fail to write all table json of %s", d.db)
	}

	return nil

}

func (d *DBInfo) GenerateSchema(w io.Writer) error {
	for _, tbl := range d.tables {
		b := append([]byte{}, []byte(tbl.CreateTableStatement(true))...)
		b = append(b, byte('\n'))

		n, err := w.Write(b)
		if err != nil {
			return errors.Errorf("Fail to write schema of %s.%s: %v", tbl.DB, tbl.Table, err)
		}
		if n != len(b) {
			return errors.Errorf("Fail to write complete schema of %s.%s", tbl.DB, tbl.Table)
		}
	}
	return nil
}

func (d *DBInfo) GenSQLite(path string) error {
	file, err := os.Create(path)
	if err != nil {
		log.Fatal("os create path", zap.String("path", path), zap.Error(err))
	}
	file.Close()

	db, err := sql.Open("sqlite3", path)
	if err != nil {
		return err
	}
	defer db.Close()

	for _, tbl := range d.tables {
		if _, err := db.Exec(tbl.CreateTableStatement(true)); err != nil {
			return err
		}

	}
	return nil
}

func GenerateSchema(db DB, outputDir string, schemaInPrompt [][]string) {
	if err := CreateDir(outputDir); err != nil {
		log.Fatal("create directory", zap.String("output dir", outputDir), zap.Error(err))
	}

	dbInfos, err := FetchAllTablesInfo(db, schemaInPrompt)
	if err != nil {
		log.Fatal("failed to fetch table info", zap.Error(err))
	}

	for db, dbInfo := range dbInfos {
		dbDir := filepath.Join(outputDir, db)
		if err := CreateDir(dbDir); err != nil {
			log.Fatal("create directory", zap.String("db dir", dbDir), zap.Error(err))
		}

		schemaFile := filepath.Join(outputDir, db, "schema.sql")

		if Exists(schemaFile) {
			log.Info("remove schema file", zap.String("schema file", schemaFile))
			if err := os.RemoveAll(schemaFile); err != nil {
				log.Fatal("Fail to remove existed schema file", zap.String("schema file", schemaFile), zap.Error(err))
			}
		}

		fSchema, err := os.Create(schemaFile)
		if err != nil {
			fSchema.Close()
			log.Fatal("Fail to create schema file", zap.String("schema file", schemaFile), zap.Error(err))
		}

		err = dbInfo.GenerateSchema(fSchema)
		fSchema.Close()
		if err != nil {
			log.Fatal("Fail to generate schema", zap.Error(err))
		}

		jsonFile := filepath.Join(outputDir, db, "tables.json")
		if Exists(jsonFile) {
			log.Info("remove existed tabels json file", zap.String("json file", jsonFile))
			if err := os.RemoveAll(jsonFile); err != nil {
				log.Fatal("Fail to remove existed tabels json file", zap.String("json file", jsonFile), zap.Error(err))
			}
		}

		fJson, err := os.Create(jsonFile)
		if err != nil {
			fJson.Close()
			log.Fatal("Fail to create tabels json file", zap.String("json file", jsonFile), zap.Error(err))
		}

		err = dbInfo.GenerateTableJson(fJson)
		fJson.Close()
		if err != nil {
			log.Fatal("Fail to generate tables json", zap.Error(err))
		}

		sqliteFile := filepath.Join(outputDir, db, fmt.Sprintf("%s.sqlite", db))
		if Exists(sqliteFile) {
			fmt.Printf("remove sqlite file %s\n", sqliteFile)
			if err := os.RemoveAll(sqliteFile); err != nil {
				log.Fatal("Fail to remove existed sqlite file", zap.String("sqlite file", sqliteFile), zap.Error(err))
			}
		}

		err = dbInfo.GenSQLite(sqliteFile)
		if err != nil {
			log.Fatal("Fail to generate sqlite", zap.String("sqlite file", sqliteFile), zap.Error(err))
		}
	}
}

func LoadSchemaFromJsonFile(file string) (*DBInfo, error) {
	fileData, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}

	data := SpiderTableJson{}

	_ = json.Unmarshal([]byte(fileData), &data)

	schemaInfo := &DBInfo{
		db:     data.DBID,
		tables: make(map[string]*TableInfo),
	}

	for _, tableName := range data.TableNamesOriginal {
		schemaInfo.tables[tableName] = &TableInfo{
			DB:      data.DBID,
			Table:   tableName,
			Columns: make(map[string]*ColumnInfo),
		}
	}

	for colIdx, colInfo := range data.ColumnsOriginal {
		tableIdx := colInfo[0].(int)
		tableName := data.TableNamesOriginal[tableIdx]
		colName := colInfo[1].(string)
		colType := data.ColumnTypes[colIdx]
		schemaInfo.tables[tableName].Columns[colName] = &ColumnInfo{
			Name: colName,
			Typ:  colType,
		}
	}

	return schemaInfo, nil
}

func FetchAllTablesInfo(db DB, tables [][]string) (map[string]*DBInfo, error) {
	dbs := make(map[string]*DBInfo)

	for _, tbl := range tables {
		dbName := tbl[0]
		tblName := tbl[1]

		tableInfo, err := db.FetchTableInfo(context.Background(), "", dbName, tblName)
		if err != nil {
			return nil, err
		}

		_, ok := dbs[dbName]
		if !ok {
			dbs[dbName] = &DBInfo{
				db:     dbName,
				tables: make(map[string]*TableInfo),
			}
		}

		dbs[dbName].tables[tblName] = tableInfo

	}

	return dbs, nil
}

func CreateDir(path string) error {
	if Exists(path) {
		if !IsDir(path) {
			return fmt.Errorf("Please make sure %s is a directory", path)
		}
		return nil
	}

	if err := os.MkdirAll(path, 0700); err != nil {
		return fmt.Errorf("Fail to create the directory %s: %v", path, err)
	}
	return nil
}

func Exists(path string) bool {
	_, err := os.Stat(path)
	if err != nil {
		return os.IsExist(err)
	}

	return true
}

func IsDir(path string) bool {
	s, err := os.Stat(path)
	if err != nil {
		return false
	}
	return s.IsDir()
}