package main

import (
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"atlas9.dev/c/demo/boot/bootdb"
)

const (
	dataModelPath        = "./docs/data_model.md"
	dataModelMermaidPath = "./docs/data_model.mmd"
)

// generateDataModel applies the migrations to a throwaway database, then
// introspects the resulting schema, so the diagram reflects what migrations
// actually produce rather than a hand-maintained description.
func generateDataModel() error {

	dir, err := os.MkdirTemp("", "demo-datamodel")
	if err != nil {
		return fmt.Errorf("creating temp dir: %w", err)
	}
	defer os.RemoveAll(dir)

	db, err := bootdb.Database(filepath.Join(dir, "schema.db"))
	if err != nil {
		return fmt.Errorf("building schema database: %w", err)
	}
	defer db.Close()

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

	mermaid := renderMermaid(tables)
	md := "# Demo Data Model\n\n" +
		"<!-- Generated by `go run ./tools`. Do not edit by hand. -->\n\n" +
		"```mermaid\n" + mermaid + "```\n"

	if err := os.MkdirAll(filepath.Dir(dataModelPath), 0o755); err != nil {
		return fmt.Errorf("creating docs dir: %w", err)
	}
	if err := os.WriteFile(dataModelPath, []byte(md), 0o644); err != nil {
		return fmt.Errorf("writing %s: %w", dataModelPath, err)
	}
	if err := os.WriteFile(dataModelMermaidPath, []byte(mermaid), 0o644); err != nil {
		return fmt.Errorf("writing %s: %w", dataModelMermaidPath, err)
	}
	return nil
}

type table struct {
	Name        string
	Columns     []column
	ForeignKeys []foreignKey
}

type column struct {
	Name     string
	Type     string
	NotNull  bool
	PK       bool
	FK       bool
	Unique   bool
}

type foreignKey struct {
	Table   string
	Columns []string
}

func loadTables(db *sql.DB) ([]table, error) {

	rows, err := db.Query(`
		SELECT name FROM sqlite_master
		WHERE type = 'table'
		AND name NOT LIKE 'sqlite_%'
		AND name != 'schema_migrations'
		ORDER BY name`)
	if err != nil {
		return nil, fmt.Errorf("listing tables: %w", err)
	}
	defer rows.Close()

	var names []string
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, fmt.Errorf("scanning table name: %w", err)
		}
		names = append(names, name)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("listing tables: %w", err)
	}

	var tables []table
	for _, name := range names {
		t, err := loadTable(db, name)
		if err != nil {
			return nil, err
		}
		tables = append(tables, t)
	}
	return tables, nil
}

func loadTable(db *sql.DB, name string) (table, error) {

	t := table{Name: name}

	fkCols := map[string]bool{}
	fks, err := loadForeignKeys(db, name, fkCols)
	if err != nil {
		return t, err
	}
	t.ForeignKeys = fks

	uniqueCols, err := loadUniqueColumns(db, name)
	if err != nil {
		return t, err
	}

	rows, err := db.Query(fmt.Sprintf("PRAGMA table_info(%q)", name))
	if err != nil {
		return t, fmt.Errorf("table_info %s: %w", name, err)
	}
	defer rows.Close()

	for rows.Next() {
		var (
			cid, notnull, pk int
			colName, colType string
			dflt             sql.NullString
		)
		if err := rows.Scan(&cid, &colName, &colType, &notnull, &dflt, &pk); err != nil {
			return t, fmt.Errorf("scanning table_info %s: %w", name, err)
		}
		t.Columns = append(t.Columns, column{
			Name:    colName,
			Type:    colType,
			NotNull: notnull != 0,
			PK:      pk != 0,
			FK:      fkCols[colName],
			Unique:  uniqueCols[colName],
		})
	}
	if err := rows.Err(); err != nil {
		return t, fmt.Errorf("table_info %s: %w", name, err)
	}
	return t, nil
}

func loadForeignKeys(db *sql.DB, name string, fkCols map[string]bool) ([]foreignKey, error) {

	rows, err := db.Query(fmt.Sprintf("PRAGMA foreign_key_list(%q)", name))
	if err != nil {
		return nil, fmt.Errorf("foreign_key_list %s: %w", name, err)
	}
	defer rows.Close()

	// Composite foreign keys appear as one row per column, grouped by id.
	byID := map[int]*foreignKey{}
	var ids []int
	for rows.Next() {
		var (
			id, seq                   int
			parent, from              string
			to                        sql.NullString
			onUpdate, onDelete, match string
		)
		if err := rows.Scan(&id, &seq, &parent, &from, &to, &onUpdate, &onDelete, &match); err != nil {
			return nil, fmt.Errorf("scanning foreign_key_list %s: %w", name, err)
		}
		fk := byID[id]
		if fk == nil {
			fk = &foreignKey{Table: parent}
			byID[id] = fk
			ids = append(ids, id)
		}
		fk.Columns = append(fk.Columns, from)
		fkCols[from] = true
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("foreign_key_list %s: %w", name, err)
	}

	sort.Ints(ids)
	var fks []foreignKey
	for _, id := range ids {
		fks = append(fks, *byID[id])
	}
	return fks, nil
}

func loadUniqueColumns(db *sql.DB, name string) (map[string]bool, error) {

	unique := map[string]bool{}

	rows, err := db.Query(fmt.Sprintf("PRAGMA index_list(%q)", name))
	if err != nil {
		return nil, fmt.Errorf("index_list %s: %w", name, err)
	}
	defer rows.Close()

	var indexes []string
	for rows.Next() {
		var (
			seq             int
			idxName, origin string
			isUnique        int
			partial         int
		)
		if err := rows.Scan(&seq, &idxName, &isUnique, &origin, &partial); err != nil {
			return nil, fmt.Errorf("scanning index_list %s: %w", name, err)
		}
		// origin "pk" is the primary key's implicit index; the PK marker
		// already covers it.
		if isUnique != 0 && partial == 0 && origin != "pk" {
			indexes = append(indexes, idxName)
		}
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("index_list %s: %w", name, err)
	}

	for _, idx := range indexes {
		cols, err := indexColumns(db, idx)
		if err != nil {
			return nil, err
		}
		// Only single-column unique indexes translate to a UK column marker.
		if len(cols) == 1 {
			unique[cols[0]] = true
		}
	}
	return unique, nil
}

func indexColumns(db *sql.DB, idx string) ([]string, error) {

	rows, err := db.Query(fmt.Sprintf("PRAGMA index_info(%q)", idx))
	if err != nil {
		return nil, fmt.Errorf("index_info %s: %w", idx, err)
	}
	defer rows.Close()

	var cols []string
	for rows.Next() {
		var seqno, cid int
		var name sql.NullString
		if err := rows.Scan(&seqno, &cid, &name); err != nil {
			return nil, fmt.Errorf("scanning index_info %s: %w", idx, err)
		}
		if name.Valid {
			cols = append(cols, name.String)
		}
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("index_info %s: %w", idx, err)
	}
	return cols, nil
}

func renderMermaid(tables []table) string {

	var b strings.Builder
	b.WriteString("erDiagram\n")

	for _, t := range tables {
		fmt.Fprintf(&b, "    %s {\n", t.Name)
		for _, c := range t.Columns {
			fmt.Fprintf(&b, "        %s %s", mermaidType(c), c.Name)
			var keys []string
			if c.PK {
				keys = append(keys, "PK")
			}
			if c.FK {
				keys = append(keys, "FK")
			}
			if c.Unique {
				keys = append(keys, "UK")
			}
			if len(keys) > 0 {
				fmt.Fprintf(&b, " %s", strings.Join(keys, ","))
			}
			if !c.NotNull && !c.PK {
				b.WriteString(` "nullable"`)
			}
			b.WriteString("\n")
		}
		b.WriteString("    }\n")
	}

	b.WriteString("\n")
	for _, t := range tables {
		for _, fk := range t.ForeignKeys {
			fmt.Fprintf(&b, "    %s }o--|| %s : %q\n", t.Name, fk.Table, strings.Join(fk.Columns, ", "))
		}
	}

	return b.String()
}

// mermaidType maps a SQLite declared type to a token mermaid accepts as an
// attribute type (no spaces or parentheses).
func mermaidType(c column) string {

	typ := strings.ToUpper(c.Type)
	if typ == "" {
		typ = "ANY"
	}
	if i := strings.IndexAny(typ, " ("); i >= 0 {
		typ = typ[:i]
	}
	return typ
}