package psql

import (
	"fmt"
	"strings"

	"github.com/go-pg/pg"
)

type TCKey struct {
	Table, Column string
}

type TTKey struct {
	Table1, Table2 string
}

type DBSchema struct {
	ColMap   map[TCKey]*DBColumn
	ColIDMap map[int]*DBColumn
	PCols    map[string]*DBColumn

	RelMap map[TTKey]*DBRel
}

type RelType int

const (
	RelBelongTo RelType = iota + 1
	RelOneToMany
	RelOneToManyThrough
)

type DBRel struct {
	Type    RelType
	Through string
	ColT    string
	Col1    string
	Col2    string
}

func NewDBSchema(db *pg.DB) (*DBSchema, error) {
	schema := &DBSchema{
		ColMap:   make(map[TCKey]*DBColumn),
		ColIDMap: make(map[int]*DBColumn),
		PCols:    make(map[string]*DBColumn),
		RelMap:   make(map[TTKey]*DBRel),
	}

	tables, err := GetTables(db)
	if err != nil {
		return nil, err
	}

	for _, t := range tables {
		cols, err := GetColumns(db, "public", t.Name)
		if err != nil {
			return nil, err
		}
		// Current table
		ct := strings.ToLower(t.Name)

		// Foreign key columns in current table
		var jcols []*DBColumn

		for _, c := range cols {
			schema.ColMap[TCKey{ct, strings.ToLower(c.Name)}] = c
			schema.ColIDMap[c.ID] = c
		}

		for _, c := range cols {
			switch {
			case c.PrimaryKey:
				schema.PCols[ct] = c

			case len(c.FKeyTable) != 0:
				if len(c.FKeyColID) == 0 {
					continue
				}

				// Foreign key column name
				ft := strings.ToLower(c.FKeyTable)
				fc, ok := schema.ColIDMap[c.FKeyColID[0]]
				if !ok {
					continue
				}

				// Belongs-to relation between current table and the
				// table in the foreign key
				rel1 := &DBRel{RelBelongTo, "", "", c.Name, fc.Name}
				schema.RelMap[TTKey{ct, ft}] = rel1

				// One-to-many relation between the foreign key table and the
				// the current table
				rel2 := &DBRel{RelOneToMany, "", "", fc.Name, c.Name}
				schema.RelMap[TTKey{ft, ct}] = rel2

				jcols = append(jcols, c)
			}
		}

		// If table contains multiple foreign key columns it's a possible
		// join table for many-to-many relationships or multiple one-to-many
		// relations

		// Below one-to-many relations use the current table as the
		// join table aka through table.
		if len(jcols) > 1 {
			col1, col2 := jcols[0], jcols[1]

			t1 := strings.ToLower(col1.FKeyTable)
			t2 := strings.ToLower(col2.FKeyTable)

			fc1, ok := schema.ColIDMap[col1.FKeyColID[0]]
			if !ok {
				continue
			}
			fc2, ok := schema.ColIDMap[col2.FKeyColID[0]]
			if !ok {
				continue
			}

			// One-to-many-through relation between 1nd foreign key table and the
			// 2nd foreign key table
			//rel1 := &DBRel{RelOneToManyThrough, ct, fc1.Name, col1.Name}
			rel1 := &DBRel{RelOneToManyThrough, ct, col2.Name, fc2.Name, col1.Name}
			schema.RelMap[TTKey{t1, t2}] = rel1

			// One-to-many-through relation between 2nd foreign key table and the
			// 1nd foreign key table
			//rel2 := &DBRel{RelOneToManyThrough, ct, col2.Name, fc2.Name}
			rel2 := &DBRel{RelOneToManyThrough, ct, col1.Name, fc1.Name, col2.Name}
			schema.RelMap[TTKey{t2, t1}] = rel2
		}
	}

	return schema, nil
}

type DBTable struct {
	Name string `sql:"name"`
	Type string `sql:"type"`
}

func GetTables(db *pg.DB) ([]*DBTable, error) {
	sqlStmt := `
	SELECT
  c.relname as "name",
  CASE c.relkind WHEN 'r' THEN 'table'
  WHEN 'v' THEN 'view'
  WHEN 'm' THEN 'materialized view'
  WHEN 'f' THEN 'foreign table' END as "type"
FROM pg_catalog.pg_class c
     LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
WHERE c.relkind IN ('r','v','m','f','')
      AND n.nspname <> 'pg_catalog'
      AND n.nspname <> 'information_schema'
      AND n.nspname !~ '^pg_toast'
  AND pg_catalog.pg_table_is_visible(c.oid);
	`

	var t []*DBTable
	_, err := db.Query(&t, sqlStmt)

	if err != nil {
		return nil, fmt.Errorf("Error fetching tables: %s", err)
	}

	return t, nil
}

type DBColumn struct {
	ID         int    `sql:"id"`
	Name       string `sql:"name"`
	Type       string `sql:"type"`
	NotNull    bool   `sql:"notnull"`
	PrimaryKey bool   `sql:"primarykey"`
	Uniquekey  bool   `sql:"uniquekey"`
	FKeyTable  string `sql:"foreignkey"`
	FKeyColID  []int  `sql:"foreignkey_fieldnum,array"`
}

func GetColumns(db *pg.DB, schema, table string) ([]*DBColumn, error) {
	sqlStmt := `
	SELECT  
    f.attnum AS id,  
    f.attname AS name,  
    f.attnotnull AS notnull,  
    pg_catalog.format_type(f.atttypid,f.atttypmod) AS type,  
    CASE  
        WHEN p.contype = 'p' THEN 't'  
        ELSE 'f'  
    END AS primarykey,  
    CASE  
        WHEN p.contype = 'u' THEN 't'  
        ELSE 'f'
    END AS uniquekey,
    CASE
        WHEN p.contype = 'f' THEN g.relname
    END AS foreignkey,
    CASE
        WHEN p.contype = 'f' THEN p.confkey
    END AS foreignkey_fieldnum
FROM pg_attribute f  
    JOIN pg_class c ON c.oid = f.attrelid  
    JOIN pg_type t ON t.oid = f.atttypid  
    LEFT JOIN pg_attrdef d ON d.adrelid = c.oid AND d.adnum = f.attnum  
    LEFT JOIN pg_namespace n ON n.oid = c.relnamespace  
    LEFT JOIN pg_constraint p ON p.conrelid = c.oid AND f.attnum = ANY (p.conkey)  
    LEFT JOIN pg_class AS g ON p.confrelid = g.oid  
WHERE c.relkind = 'r'::char  
    AND n.nspname = $1  -- Replace with Schema name  
    AND c.relname = $2  -- Replace with table name  
    AND f.attnum > 0 ORDER BY id;
	`

	stmt, err := db.Prepare(sqlStmt)
	if err != nil {
		return nil, fmt.Errorf("Error fetching columns: %s", err)
	}

	var t []*DBColumn
	_, err = stmt.Query(&t, schema, table)
	if err != nil {
		return nil, fmt.Errorf("Error fetching columns: %s", err)
	}

	return t, nil
}