package main

import (
	// "database/sql"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/gdamore/tcell"
	"github.com/jmoiron/sqlx"
	_ "github.com/lib/pq"
	"github.com/rivo/tview"
)

var app = tview.NewApplication()

func main() {
	db, err := sqlx.Open("postgres", "postgres://postgres:password@localhost:5432/postgres?sslmode=disable")
	if err != nil {
		log.Printf("could not connect to db: %v", err)
		os.Exit(1)
	}

	schemata := newSchemaSelector(db)
	queries := newQuerySelector(db)
	pv := newPlanView()
	flex := tview.NewFlex().
		AddItem(schemata, 30, 0, true).
		AddItem(tview.NewFlex().
			SetDirection(tview.FlexRow).
			AddItem(queries, 0, 25, false).
			AddItem(pv, 0, 75, false), 0, 100, false)

	schemata.SetSelectedFunc(func(_ int, t, _ string, _ rune) {
		go func() {
			app.QueueUpdateDraw(func() {
				if queries.SelectSchema(t) {
					app.SetFocus(queries)
				}
			})
		}()
	})

	app.SetInputCapture(func(event *tcell.EventKey) *tcell.EventKey {
		switch event.Key() {
		case tcell.KeyF1:
			app.SetFocus(schemata)
		case tcell.KeyF2:
			app.SetFocus(queries)
		case tcell.KeyF3:
			app.SetFocus(pv)
		case tcell.KeyEnter:
			if app.GetFocus() == queries {
				p, err := queries.Analyze()
				if err != nil {
					pv.SetError(err)
				} else {
					pv.SetPlan(p)
					app.SetFocus(pv)
				}
			}
		case tcell.KeyTAB:
			if app.GetFocus() == pv.tree {
				app.SetFocus(pv.detail)
			} else if app.GetFocus() == pv.detail {
				app.SetFocus(pv.tree)
			}
		}
		return event
	})

	if err := app.SetRoot(flex, true).Run(); err != nil {
		log.Printf("could not create application: %v", err)
		os.Exit(1)
	}
}

func newSchemaSelector(db *sqlx.DB) *tview.List {
	v := tview.NewList()
	v.SetTitle("Schemata")
	v.SetBorder(true)
	v.ShowSecondaryText(false)
	v.SetHighlightFullLine(true)

	// v.SetMainTextColor(tcell.ColorDefault)
	// v.SetBackgroundColor(tcell.ColorDefault)
	// v.SetSelectedTextColor(tcell.ColorRed)
	// v.SetSelectedBackgroundColor(tcell.ColorDefault)

	type schema struct {
		SchemaName string `db:"schema_name"`
	}

	var schemata []*schema
	if err := db.Select(&schemata, "SELECT schema_name FROM information_schema.schemata"); err != nil {
		v.AddItem(fmt.Sprintf("could not find schemata: %v", err), "", 'a', nil)
		return v
	}

	for _, s := range schemata {
		v.AddItem(s.SchemaName, "", 0, nil)
	}

	return v
}

type pg_stat_statement struct {
	UserID            int     `db:"userid"`
	DBID              int     `db:"dbid"`
	QueryID           int64   `db:"queryid"`
	Query             string  `db:"query"`
	Calls             int     `db:"calls"`
	TotalTime         float32 `db:"total_time"`
	MinTime           float32 `db:"min_time"`
	MaxTime           float32 `db:"max_time"`
	MeanTime          float32 `db:"mean_time"`
	StddevTime        float32 `db:"stddev_time"`
	Rows              int64   `db:"rows"`
	SharedBlksHit     int64   `db:"shared_blks_hit"`
	SharedBlksRead    int64   `db:"shared_blks_read"`
	SharedBlksDirtied int64   `db:"shared_blks_dirtied"`
	SharedBlksWritten int64   `db:"shared_blks_written"`
	LocalBlksHit      int64   `db:"local_blks_hit"`
	LocalBlksRead     int64   `db:"local_blks_read"`
	LocalBlksDirtied  int64   `db:"local_blks_dirtied"`
	LocalBlksWritten  int64   `db:"local_blks_written"`
	TempBlksRead      int64   `db:"temp_blks_read"`
	TempBlksWritten   int64   `db:"temp_blks_written"`
	BlkReadTime       float32 `db:"blk_read_time"`
	BlkWriteTime      float32 `db:"blk_write_time"`
}

type querySelector struct {
	*tview.Table
	db    *sqlx.DB
	stmts []*pg_stat_statement
}

func newQuerySelector(db *sqlx.DB) *querySelector {
	v := tview.NewTable()
	v.SetTitle("Queries")
	v.SetBorder(true)
	v.SetSelectable(true, false)
	// v.SetBackgroundColor(tcell.ColorDefault)
	// v.SetSelectedStyle(tcell.ColorDefault, tcell.ColorDefault, tcell.AttrReverse)

	// v.SetMainTextColor(tcell.ColorDefault)
	// v.SetBackgroundColor(tcell.ColorDefault)
	// v.SetSelectedTextColor(tcell.ColorRed)
	// v.SetSelectedBackgroundColor(tcell.ColorDefault)

	qs := &querySelector{
		Table: v,
		db:    db,
	}

	return qs
}

func (q *querySelector) SelectSchema(schema string) bool {
	q.Clear()

	q.stmts = make([]*pg_stat_statement, 0)

	if err := q.db.Select(&q.stmts, fmt.Sprintf("SELECT * FROM %s.pg_stat_statements ORDER BY total_time DESC", schema)); err != nil {
		log.Printf("could not select schema: %v", err)
		cell := tview.NewTableCell(fmt.Sprintf("could not find queries: %v", err))
		// cell.SetTextColor(tcell.ColorYellow)
		cell.SetAlign(tview.AlignCenter)
		cell.SetSelectable(false)
		q.SetCell(0, 0, cell)
		return false
	}

	q.SetFixed(1, 0)

	newHeaderCell := func(t string) *tview.TableCell {
		cell := tview.NewTableCell(t)
		cell.SetSelectable(false)
		// cell.SetTextColor(tcell.ColorGreen)
		return cell
	}

	q.SetCell(0, 0, newHeaderCell("User ID"))
	q.SetCell(0, 1, newHeaderCell("Query"))
	q.SetCell(0, 2, newHeaderCell("Calls"))
	q.SetCell(0, 3, newHeaderCell("Total Time"))

	newTableCell := func(t string) *tview.TableCell {
		cell := tview.NewTableCell(t)
		// cell.SetStyle(tcell.StyleDefault)
		return cell
	}

	for idx, stmt := range q.stmts {
		q.SetCell(idx+1, 0, newTableCell(strconv.Itoa(stmt.UserID)))
		q.SetCell(idx+1, 1, newTableCell(stmt.Query))
		q.SetCell(idx+1, 2, newTableCell(strconv.Itoa(stmt.Calls)))
		q.SetCell(idx+1, 3, newTableCell(strconv.Itoa(int(stmt.TotalTime))))
	}

	return true
}

func (q *querySelector) Analyze() ([]*analysis, error) {
	row, _ := q.GetSelection()
	if len(q.stmts) < row {
		return nil, fmt.Errorf("no queries to analyze")
	}
	stmt := q.stmts[row-1]
	tx, err := q.db.Beginx()
	if err != nil {
		return nil, fmt.Errorf("could not begin transaction: %v", err)
	}
	defer tx.Rollback()

	type qp struct {
		Plan string `db:"QUERY PLAN"`
	}

	var plan qp
	if err := tx.Get(&plan, fmt.Sprintf("EXPLAIN (ANALYZE, COSTS, BUFFERS, VERBOSE, FORMAT JSON) %s", stmt.Query)); err != nil {
		return nil, fmt.Errorf("could not analyze query: %v", err)
	}

	if err := tx.Rollback(); err != nil {
		return nil, fmt.Errorf("could not roll back transaction: %v", err)
	}

	a := make([]*analysis, 0)
	// if err := json.Unmarshal([]byte(plan.Plan), &a); err != nil {
	if err := json.Unmarshal([]byte(testPlan), &a); err != nil {
		return nil, fmt.Errorf("could not unmarshal plan: %v", err)
	}

	return a, nil
}

type analysis struct {
	Plan         *plan   `json:"Plan"`
	PlanningTime float32 `json:"Planning Time"`
	// Triggers      []string `json:"Triggers"`
	ExecutionTime float32 `json:"Execution Time"`
}

type plan struct {
	NodeType                  string   `json:"Node Type"`
	StartupCost               float32  `json:"Startup Cost"`
	TotalCost                 float32  `json:"Total Cost"`
	PlanRows                  int      `json:"Plan Rows"`
	PlanWidth                 int      `json:"Plan Width"`
	ActualStartupTime         float32  `json:"Actual Startup Time"`
	ActualTotalTime           float32  `json:"Actual Total Time"`
	ActualRows                int      `json:"Actual Rows"`
	ActualLoops               int      `json:"Actual Loops"`
	Output                    []string `json:"Output"`
	SharedHitBlocks           int      `json:"Shared Hit Blocks"`
	SharedReadBlocks          int      `json:"Shared Read Blocks"`
	SharedDirtiedBlocks       int      `json:"Shared Dirtied Blocks"`
	SharedWrittenBlocks       int      `json:"Shared Written Blocks"`
	LocalHitBlocks            int      `json:"Local Hit Blocks"`
	LocalReadBlocks           int      `json:"Local Read Blocks"`
	LocalDirtiedBlocks        int      `json:"Local Dirtied Blocks"`
	LocalWrittenBlocks        int      `json:"Local Written Blocks"`
	TempReadBlocks            int      `json:"Temp Read Blocks"`
	TempWrittenBlocks         int      `json:"Temp Written Blocks"`
	IOReadTime                float32  `json:"I/O Read Time"`
	IOWriteTime               float32  `json:"I/O Write Time"`
	ParentRelationship        *string  `json:"Parent Relationship"`
	SortKey                   []string `json:"Sort Key"`
	SortMethod                *string  `json:"Sort Method"`
	SortSpaceUsed             *int     `json:"Sort Space Used"`
	SortSpaceType             *string  `json:"Sort Space Type"`
	JoinType                  *string  `json:"Join Type"`
	Strategy                  *string  `json:"Strategy"`
	RelationName              *string  `json:"Relation Name"`
	Schema                    *string  `json:"Schema"`
	Alias                     *string  `json:"Alias"`
	RecheckCond               *string  `json:"Recheck Cond"`
	RowsRemovedByIndexRecheck *int     `json:"Rows Removed by Index Recheck"`
	ExactHeapBlocks           *int     `json:"Exact Heap Blocks"`
	LossyHeapBlocks           *int     `json:"Lossy Heap Blocks"`
	IndexName                 *string  `json:"Index Name"`
	IndexCond                 *string  `json:"Index Cond"`
	ScanDirection             *string  `json:"Scan Direction"`
	ParallelAware             *bool    `json:"Parallel Aware"`
	FunctionName              *string  `json:"Function Name"`
	Plans                     []*plan  `json:"Plans"`
}

type planView struct {
	*tview.Flex

	tree   *tview.TreeView
	detail *tview.Table
}

func newPlanView() *planView {
	f := tview.NewFlex()
	f.SetBorder(true).SetTitle("Plan")
	f.SetDirection(tview.FlexRow)

	t := tview.NewTreeView()
	d := tview.NewTable()

	f.AddItem(t, 0, 50, true)
	f.AddItem(d, 0, 50, false)

	pv := &planView{
		Flex:   f,
		tree:   t,
		detail: d,
	}

	t.SetSelectedFunc(func(node *tview.TreeNode) {
		go func() {
			app.QueueUpdateDraw(func() {
				ref := node.GetReference()
				if ref == nil {
					pv.ShowDetail(nil)
					return
				}
				p, ok := ref.(*plan)
				if !ok {
					pv.ShowDetail(nil)
					return
				}
				log.Printf("selected node named %s", p.NodeType)
				pv.ShowDetail(p)
			})
		}()
	})
	return pv
}

func (p *planView) SetError(err error) {
	p.tree.SetRoot(tview.NewTreeNode(err.Error()))
}

func (p *planView) SetPlan(as []*analysis) {
	if len(as) == 0 {
		p.tree.SetRoot(tview.NewTreeNode("No plan results to display"))
		return
	}

	if len(as) > 1 {
		log.Printf("skipping display of more than one analysis result")
	}

	a := as[0]
	root := tview.NewTreeNode(fmt.Sprintf("Query Plan (planning=%s execution=%s)", f2d(a.PlanningTime), f2d(a.ExecutionTime)))
	root.AddChild(buildNode(a.Plan))
	p.tree.SetRoot(root)
	p.tree.SetCurrentNode(root)
}

func buildNode(p *plan) *tview.TreeNode {
	dur := f2d(p.ActualTotalTime)
	n := tview.NewTreeNode(fmt.Sprintf("%s (%s)", p.NodeType, dur.String()))
	n.SetSelectable(true)
	n.SetReference(p)
	for _, subplan := range p.Plans {
		n.AddChild(buildNode(subplan))
	}
	return n
}

func (pv *planView) ShowDetail(p *plan) {
	pv.detail.Clear()

	if p == nil {
		return
	}

	pv.detail.SetCellSimple(0, 0, "Node Type")
	pv.detail.SetCellSimple(0, 1, p.NodeType)

	pv.detail.SetCellSimple(1, 0, "Startup Cost")
	pv.detail.SetCellSimple(1, 1, f2a(p.StartupCost))

	pv.detail.SetCellSimple(2, 0, "Total Cost")
	pv.detail.SetCellSimple(2, 1, f2a(p.TotalCost))

	pv.detail.SetCellSimple(3, 0, "Plan Rows")
	pv.detail.SetCellSimple(3, 1, strconv.Itoa(p.PlanRows))

	pv.detail.SetCellSimple(4, 0, "Actual Rows")
	pv.detail.SetCellSimple(4, 1, strconv.Itoa(p.ActualRows))

	pv.detail.SetCellSimple(5, 0, "Plan Width")
	pv.detail.SetCellSimple(5, 1, strconv.Itoa(p.PlanWidth))

	pv.detail.SetCellSimple(6, 0, "Actual Startup Time")
	pv.detail.SetCellSimple(6, 1, f2d(p.ActualStartupTime).String())

	pv.detail.SetCellSimple(7, 0, "Actual Total Time")
	pv.detail.SetCellSimple(7, 1, f2d(p.ActualTotalTime).String())

	pv.detail.SetCellSimple(8, 0, "Actual Loops")
	pv.detail.SetCellSimple(8, 1, strconv.Itoa(p.ActualLoops))

	pv.detail.SetCellSimple(9, 0, "Shared Hit Blocks")
	pv.detail.SetCellSimple(9, 1, strconv.Itoa(p.SharedHitBlocks))

	pv.detail.SetCellSimple(10, 0, "Shared Read Blocks")
	pv.detail.SetCellSimple(10, 1, strconv.Itoa(p.SharedReadBlocks))

	pv.detail.SetCellSimple(11, 0, "Shared Dirtied Blocks")
	pv.detail.SetCellSimple(11, 1, strconv.Itoa(p.SharedDirtiedBlocks))

	pv.detail.SetCellSimple(12, 0, "Shared Written Blocks")
	pv.detail.SetCellSimple(12, 1, strconv.Itoa(p.SharedWrittenBlocks))

	pv.detail.SetCellSimple(13, 0, "Local Hit Blocks")
	pv.detail.SetCellSimple(13, 1, strconv.Itoa(p.LocalHitBlocks))

	pv.detail.SetCellSimple(14, 0, "Local Read Blocks")
	pv.detail.SetCellSimple(14, 1, strconv.Itoa(p.LocalReadBlocks))

	pv.detail.SetCellSimple(15, 0, "Local Dirtied Blocks")
	pv.detail.SetCellSimple(15, 1, strconv.Itoa(p.LocalDirtiedBlocks))

	pv.detail.SetCellSimple(16, 0, "Local Written Blocks")
	pv.detail.SetCellSimple(16, 1, strconv.Itoa(p.LocalWrittenBlocks))

	pv.detail.SetCellSimple(17, 0, "Temp Read Blocks")
	pv.detail.SetCellSimple(17, 1, strconv.Itoa(p.TempReadBlocks))

	pv.detail.SetCellSimple(18, 0, "Temp Written Blocks")
	pv.detail.SetCellSimple(18, 1, strconv.Itoa(p.TempWrittenBlocks))

	pv.detail.SetCellSimple(19, 0, "I/O Read Time")
	pv.detail.SetCellSimple(19, 1, f2d(p.IOReadTime).String())

	pv.detail.SetCellSimple(20, 0, "I/O Write Time")
	pv.detail.SetCellSimple(20, 1, f2d(p.IOWriteTime).String())

	idx := 21

	if p.ParentRelationship != nil {
		pv.detail.SetCellSimple(idx, 0, "Parent Relationship")
		pv.detail.SetCellSimple(idx, 1, *p.ParentRelationship)
		idx++
	}

	if len(p.Output) > 0 {
		pv.detail.SetCellSimple(idx, 0, "Output")
		pv.detail.SetCellSimple(idx, 1, strings.Join(p.Output, ", "))
		idx++
	}

	if len(p.SortKey) > 0 {
		pv.detail.SetCellSimple(idx, 0, "Sort Key")
		pv.detail.SetCellSimple(idx, 1, strings.Join(p.SortKey, ", "))
		idx++
	}

	if p.SortMethod != nil {
		pv.detail.SetCellSimple(idx, 0, "Sort Method")
		pv.detail.SetCellSimple(idx, 1, *p.SortMethod)
		idx++
	}

	if p.SortSpaceUsed != nil {
		pv.detail.SetCellSimple(idx, 0, "Sort Space Used")
		pv.detail.SetCellSimple(idx, 1, strconv.Itoa(*p.SortSpaceUsed))
		idx++
	}

	if p.SortSpaceType != nil {
		pv.detail.SetCellSimple(idx, 0, "Sort Space Type")
		pv.detail.SetCellSimple(idx, 1, *p.SortSpaceType)
		idx++
	}

	if p.JoinType != nil {
		pv.detail.SetCellSimple(idx, 0, "Join Type")
		pv.detail.SetCellSimple(idx, 1, *p.JoinType)
		idx++
	}

	if p.Strategy != nil {
		pv.detail.SetCellSimple(idx, 0, "Strategy")
		pv.detail.SetCellSimple(idx, 1, *p.Strategy)
		idx++
	}

	if p.RelationName != nil {
		pv.detail.SetCellSimple(idx, 0, "Relation Name")
		pv.detail.SetCellSimple(idx, 1, *p.RelationName)
		idx++
	}

	if p.Schema != nil {
		pv.detail.SetCellSimple(idx, 0, "Schema")
		pv.detail.SetCellSimple(idx, 1, *p.Schema)
		idx++
	}

	if p.Alias != nil {
		pv.detail.SetCellSimple(idx, 0, "Alias")
		pv.detail.SetCellSimple(idx, 1, *p.Alias)
		idx++
	}

	if p.RecheckCond != nil {
		pv.detail.SetCellSimple(idx, 0, "Recheck Cond")
		pv.detail.SetCellSimple(idx, 1, *p.RecheckCond)
		idx++
	}

	if p.RowsRemovedByIndexRecheck != nil {
		pv.detail.SetCellSimple(idx, 0, "Rows Removed by Index Recheck")
		pv.detail.SetCellSimple(idx, 1, strconv.Itoa(*p.RowsRemovedByIndexRecheck))
		idx++
	}

	if p.ExactHeapBlocks != nil {
		pv.detail.SetCellSimple(idx, 0, "Exact Heap Blocks")
		pv.detail.SetCellSimple(idx, 1, strconv.Itoa(*p.ExactHeapBlocks))
		idx++
	}

	if p.LossyHeapBlocks != nil {
		pv.detail.SetCellSimple(idx, 0, "Lossy Heap Blocks")
		pv.detail.SetCellSimple(idx, 1, strconv.Itoa(*p.LossyHeapBlocks))
		idx++
	}

	if p.IndexName != nil {
		pv.detail.SetCellSimple(idx, 0, "Index Name")
		pv.detail.SetCellSimple(idx, 1, *p.IndexName)
		idx++
	}

	if p.IndexCond != nil {
		pv.detail.SetCellSimple(idx, 0, "Index Cond")
		pv.detail.SetCellSimple(idx, 1, *p.IndexCond)
		idx++
	}

	if p.ScanDirection != nil {
		pv.detail.SetCellSimple(idx, 0, "Scan Direction")
		pv.detail.SetCellSimple(idx, 1, *p.ScanDirection)
		idx++
	}

	if p.ParallelAware != nil {
		pv.detail.SetCellSimple(idx, 0, "Parallel Aware")
		pv.detail.SetCellSimple(idx, 1, fmt.Sprintf("%t", *p.ParallelAware))
		idx++
	}

	if p.FunctionName != nil {
		pv.detail.SetCellSimple(idx, 0, "Function Name")
		pv.detail.SetCellSimple(idx, 1, *p.FunctionName)
		idx++
	}
}

func f2a(f float32) string {
	return fmt.Sprintf("%f", f)
}

func f2d(f float32) time.Duration {
	return time.Duration(int(f)) * time.Millisecond
	// ms := time.Duration(int(f)) * time.Millisecond
	// ns := time.Duration(int((f-float32(ms))*1000)) * time.Nanosecond
	// return time.Duration(ms + ns)
}