package oas_test

import (
	"context"
	"dataservices/commons/oas"
	"testing"

	"github.com/tj/assert"
)

func NewTestExtractor() *oas.Extractor {
	metaColumns := []*oas.MetaColumn{
		{
			TableSchema: "db1",
			TableName:   "t1",
			ColumnName:  "a",
			DataType:    "varchar",
			IsNullable:  "YES",
		},
		{
			TableSchema: "db1",
			TableName:   "t1",
			ColumnName:  "b",
			DataType:    "varchar",
			IsNullable:  "YES",
		},
		{
			TableSchema: "db1",
			TableName:   "t2",
			ColumnName:  "c",
			DataType:    "varchar",
			IsNullable:  "YES",
		},
		{
			TableSchema: "db1",
			TableName:   "t2",
			ColumnName:  "d",
			DataType:    "varchar",
			IsNullable:  "YES",
		},
		{
			TableSchema: "db2",
			TableName:   "t1",
			ColumnName:  "e",
			DataType:    "varchar",
			IsNullable:  "YES",
		},
		{
			TableSchema: "db2",
			TableName:   "t1",
			ColumnName:  "f",
			DataType:    "varchar",
			IsNullable:  "YES",
		},
		{
			TableSchema: "db2",
			TableName:   "t2",
			ColumnName:  "g",
			DataType:    "varchar",
			IsNullable:  "YES",
		},
		{
			TableSchema: "db2",
			TableName:   "t2",
			ColumnName:  "h",
			DataType:    "varchar",
			IsNullable:  "YES",
		},
	}
	return oas.NewExtractor(metaColumns)
}

func TestBasicSelect(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "select * from db1.t1")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "b", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "select * from db1.t2")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "c", columns[0].ColumnName)
	assert.Equal(t, "d", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "select t.* from db1.t2 as t")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "c", columns[0].ColumnName)
	assert.Equal(t, "d", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "select * from db2.t1")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "e", columns[0].ColumnName)
	assert.Equal(t, "f", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "select * from db2.t2")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "g", columns[0].ColumnName)
	assert.Equal(t, "h", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "select t2.g from db2.t2")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "g", columns[0].ColumnName)
	columns, err = extractor.Extract(ctx, "select *, rank() over (order by g) as t from db2.t2")
	assert.Nil(t, err)
	assert.Equal(t, 3, len(columns))
	assert.Equal(t, "g", columns[0].ColumnName)
	assert.Equal(t, "h", columns[1].ColumnName)
	assert.Equal(t, "t", columns[2].ColumnName)
}

func TestNoDatabase(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	_, err := extractor.Extract(ctx, "select * from t1")
	assert.NotNil(t, err)
}

func TestUse(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "use db1; select * from t1")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "b", columns[1].ColumnName)
}

func TestBasicJoin(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "use db1; select * from t1 join t2 on t1.a=t2.c;")
	assert.Nil(t, err)
	assert.Equal(t, 4, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "b", columns[1].ColumnName)
	assert.Equal(t, "c", columns[2].ColumnName)
	assert.Equal(t, "d", columns[3].ColumnName)
	columns, err = extractor.Extract(ctx, "use db1; select t1.*, t2.* from t1, t2;")
	assert.Nil(t, err)
	assert.Equal(t, 4, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "b", columns[1].ColumnName)
	assert.Equal(t, "c", columns[2].ColumnName)
	assert.Equal(t, "d", columns[3].ColumnName)
	columns, err = extractor.Extract(ctx, "use db1; select t1.* from t1, t2;")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "b", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "use db1; select t1.* from t1, t2;")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "b", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "select t1.* from db1.t1, db2.t1;")
	assert.Nil(t, err)
	assert.Equal(t, 4, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "b", columns[1].ColumnName)
	assert.Equal(t, "e", columns[2].ColumnName)
	assert.Equal(t, "f", columns[3].ColumnName)
	columns, err = extractor.Extract(ctx, "select db1.t1.* from db1.t1, db2.t1;")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "b", columns[1].ColumnName)
}

func TestProjection(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "select b from db1.t1")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "b", columns[0].ColumnName)
	columns, err = extractor.Extract(ctx, "select a as a1 from db1.t1")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "a1", columns[0].ColumnName)
	columns, err = extractor.Extract(ctx, "select a a2 from db1.t1")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "a2", columns[0].ColumnName)
	columns, err = extractor.Extract(ctx, "select max(b) from db1.t1")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "max(b)", columns[0].ColumnName)
}

func TestSelectConst(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "select \"Hello World\"")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "Hello World", columns[0].ColumnName)
	columns, err = extractor.Extract(ctx, "select 1")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "1", columns[0].ColumnName)
	columns, err = extractor.Extract(ctx, "select 3.14")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "3.14", columns[0].ColumnName)
}

func TestWindowFunction(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "select a, count(*) over (partition by a) as c from db1.t1")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "c", columns[1].ColumnName)
}

func TestAggregation(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "select a from db1.t1 group by a")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	columns, err = extractor.Extract(ctx, "select t.a, sum(b) from db1.t1 group by a")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "sum(b)", columns[1].ColumnName)
}

func TestSubQuery(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "select * from (select a, count(*) over (partition by a) as c from db1.t1) tt1")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "c", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "select tt1.* from (select a, count(*) over (partition by a) as c from db1.t1) tt1")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "c", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "select TT1.* from (select a, count(*) over (partition by a) as c from db1.t1) tt1")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "a", columns[0].ColumnName)
	assert.Equal(t, "c", columns[1].ColumnName)
}

func TestUnion(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "select 1, 2 union select 3, 4")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "1", columns[0].ColumnName)
	assert.Equal(t, "2", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "select 1, 2 union select 3, 4 intersect select 3,4 EXCEPT select 1,2;")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "1", columns[0].ColumnName)
	assert.Equal(t, "2", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "select * from (select 1, 2 union select 3, 4 intersect select 3,4 EXCEPT select 1,2);")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "1", columns[0].ColumnName)
	assert.Equal(t, "2", columns[1].ColumnName)
}

func TestBasicCTE(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "with cte(i,j) as (select 1,2 UNION select 3,4) select * from cte;")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "i", columns[0].ColumnName)
	assert.Equal(t, "j", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "with cte as (select 1,2 UNION select 3,4) select * from cte;")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "1", columns[0].ColumnName)
	assert.Equal(t, "2", columns[1].ColumnName)
	columns, err = extractor.Extract(ctx, "with cte(m,n) as (with cte(i,j) as (select 1,2 UNION select 3,4) select * from cte) select * from cte;")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "m", columns[0].ColumnName)
	assert.Equal(t, "n", columns[1].ColumnName)
}

func TestRecursiveCTE(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "WITH RECURSIVE cte (n) AS ( SELECT 1 UNION ALL SELECT n + 1 FROM cte WHERE n < 5) SELECT * FROM cte;")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "n", columns[0].ColumnName)
	columns, err = extractor.Extract(ctx, "WITH RECURSIVE cte AS ( SELECT 1 AS n, 'abc' AS str UNION ALL SELECT n + 1, CONCAT(str, str) FROM cte WHERE n < 3) SELECT * FROM cte;")
	assert.Nil(t, err)
	assert.Equal(t, 2, len(columns))
	assert.Equal(t, "n", columns[0].ColumnName)
	assert.Equal(t, "str", columns[1].ColumnName)
}

func TestShow(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "show tables in db1;")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "Tables_in_db1", columns[0].ColumnName)
	columns, err = extractor.Extract(ctx, "use db1; show tables")
	assert.Nil(t, err)
	assert.Equal(t, 1, len(columns))
	assert.Equal(t, "Tables_in_db1", columns[0].ColumnName)
}

func TestExplain(t *testing.T) {
	extractor := NewTestExtractor()
	ctx := context.Background()
	columns, err := extractor.Extract(ctx, "explain select * from db1.t1")
	assert.Nil(t, err)
	assert.Equal(t, 5, len(columns))
}