package store

import (
	"axolotl/models"
	"database/sql"
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"strings"
	"time"

	_ "modernc.org/sqlite"
)

var migrations = []string{
	`CREATE TABLE IF NOT EXISTS nodes (id TEXT PRIMARY KEY, title TEXT NOT NULL, content TEXT, due_date DATETIME, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP)`,
	`CREATE TABLE IF NOT EXISTS rels (from_id TEXT NOT NULL, rel_name TEXT NOT NULL, to_id TEXT NOT NULL DEFAULT '', PRIMARY KEY (from_id, rel_name, to_id), FOREIGN KEY (from_id) REFERENCES nodes(id) ON DELETE CASCADE)`,
	`CREATE INDEX IF NOT EXISTS idx_rels_from ON rels(from_id)`,
	`CREATE INDEX IF NOT EXISTS idx_rels_to ON rels(to_id)`,
}

// querier abstracts *sql.DB and *sql.Tx so SQL helpers work for both.
type querier interface {
	Exec(query string, args ...any) (sql.Result, error)
	Query(query string, args ...any) (*sql.Rows, error)
	QueryRow(query string, args ...any) *sql.Row
}

// GraphStoreSqlite is the top-level Store backed by a SQLite database file.
type GraphStoreSqlite struct {
	db *sql.DB
}

// txStore wraps an active transaction. Its Transaction method is a no-op
// passthrough so nested calls reuse the same transaction.
type txStore struct {
	db *sql.DB
	tx *sql.Tx
}

// InitSQLiteStore creates the database file and applies the schema.
// It is idempotent (uses CREATE TABLE IF NOT EXISTS).
func InitSQLiteStore(path string) error {
	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return err
	}
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return err
	}
	defer db.Close()
	pragmas := []string{"PRAGMA journal_mode=WAL", "PRAGMA busy_timeout=5000", "PRAGMA foreign_keys=ON"}
	for _, q := range append(pragmas, migrations...) {
		if _, err := db.Exec(q); err != nil {
			return err
		}
	}
	return nil
}

// FindAndOpenSQLiteStore opens the SQLite database by walking up from the
// current working directory to find an .ax/ax.db file.
func FindAndOpenSQLiteStore() (GraphStore, error) {
	dataRoot, err := FindDataRoot(".local", "share")
	if err != nil {
		return nil, fmt.Errorf("failed to find data dir: %w", err)
	}
	dbPath := filepath.Join(dataRoot, "ax.db")
	if _, err := os.Stat(dbPath); err != nil {
		return nil, fmt.Errorf("no database found at %s: run 'ax init' first", dbPath)
	}
	return NewSQLiteStore(dbPath)
}

// FindOrInitSQLiteStore is like FindAndOpenSQLiteStore but intended for server
// mode: if no database is found it creates and initialises one in the
// ~/.local/share/ax/ directory instead of returning an error.
func FindOrInitSQLiteStore() (GraphStore, error) {
	dataRoot, err := FindDataRoot(".local", "share")
	if err != nil {
		return nil, fmt.Errorf("failed to find data dir: %w", err)
	}
	dbPath := filepath.Join(dataRoot, "ax.db")
	if _, err := os.Stat(dbPath); err != nil {
		if err := InitSQLiteStore(dbPath); err != nil {
			return nil, err
		}
	}
	return NewSQLiteStore(dbPath)
}

// NewSQLiteStore opens a SQLite database at the given path, runs a one-time
// schema migration if needed, then applies per-connection PRAGMAs.
func NewSQLiteStore(path string) (GraphStore, error) {
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}
	// FK must be OFF during migration (table drops/renames).
	for _, q := range []string{"PRAGMA journal_mode=WAL", "PRAGMA busy_timeout=5000", "PRAGMA foreign_keys=OFF"} {
		if _, err := db.Exec(q); err != nil {
			db.Close()
			return nil, err
		}
	}
	if err := migrateSchema(db); err != nil {
		db.Close()
		return nil, fmt.Errorf("schema migration failed: %w", err)
	}
	if err := migrateDueDateColumn(db); err != nil {
		db.Close()
		return nil, fmt.Errorf("due_date column migration failed: %w", err)
	}
	if _, err := db.Exec("PRAGMA foreign_keys=ON"); err != nil {
		db.Close()
		return nil, err
	}
	return &GraphStoreSqlite{db: db}, nil
}

// migrateSchema migrates from the legacy two-table (tags + rels) schema to the
// unified rels-only schema. It is a no-op when migration has already been applied.
func migrateSchema(db *sql.DB) error {
	var tagsExists int
	if err := db.QueryRow("SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='tags'").Scan(&tagsExists); err != nil {
		return err
	}
	if tagsExists == 0 {
		return nil // already on new schema
	}
	for _, stmt := range []string{
		`CREATE TABLE rels_new (from_id TEXT NOT NULL, rel_name TEXT NOT NULL, to_id TEXT NOT NULL DEFAULT '', PRIMARY KEY (from_id, rel_name, to_id), FOREIGN KEY (from_id) REFERENCES nodes(id) ON DELETE CASCADE)`,
		`INSERT OR IGNORE INTO rels_new (from_id, rel_name, to_id) SELECT from_id, rel_type, to_id FROM rels`,
		`INSERT OR IGNORE INTO rels_new (from_id, rel_name, to_id) SELECT node_id, tag, '' FROM tags`,
		`DROP TABLE rels`,
		`DROP TABLE tags`,
		`ALTER TABLE rels_new RENAME TO rels`,
		`CREATE INDEX IF NOT EXISTS idx_rels_from ON rels(from_id)`,
		`CREATE INDEX IF NOT EXISTS idx_rels_to ON rels(to_id)`,
	} {
		if _, err := db.Exec(stmt); err != nil {
			return err
		}
	}
	return nil
}

// migrateDueDateColumn converts the due_date column from TEXT to DATETIME affinity
// for databases created before this schema change. It is a no-op when already migrated.
func migrateDueDateColumn(db *sql.DB) error {
	rows, err := db.Query("PRAGMA table_info(nodes)")
	if err != nil {
		return err
	}
	defer rows.Close()
	var colType string
	for rows.Next() {
		var cid, notNull, pk int
		var name, typ string
		var dfltVal any
		if err := rows.Scan(&cid, &name, &typ, &notNull, &dfltVal, &pk); err != nil {
			return err
		}
		if name == "due_date" {
			colType = strings.ToUpper(typ)
			break
		}
	}
	rows.Close()
	if colType == "DATETIME" || colType == "" {
		return nil // already on new schema or column missing
	}
	for _, stmt := range []string{
		`CREATE TABLE nodes_new (id TEXT PRIMARY KEY, title TEXT NOT NULL, content TEXT, due_date DATETIME, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP)`,
		`INSERT INTO nodes_new SELECT id, title, content, due_date, created_at, updated_at FROM nodes`,
		`DROP TABLE nodes`,
		`ALTER TABLE nodes_new RENAME TO nodes`,
	} {
		if _, err := db.Exec(stmt); err != nil {
			return err
		}
	}
	return nil
}

// --- Transaction ---

func (s *GraphStoreSqlite) Transaction(fn func(GraphStore) error) error {
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if err := fn(&txStore{db: s.db, tx: tx}); err != nil {
		return err
	}
	return tx.Commit()
}

func (s *txStore) Transaction(fn func(GraphStore) error) error {
	return fn(s) // already in a transaction — reuse it
}

// --- Node operations ---

func addNode(q querier, id, title, content string, dueDate *time.Time, createdAt, updatedAt string) error {
	var dd any
	if dueDate != nil {
		dd = dueDate.UTC().Format("2006-01-02")
	}
	_, err := q.Exec(
		"INSERT INTO nodes (id, title, content, due_date, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)",
		id, title, content, dd, createdAt, updatedAt,
	)
	return err
}

func getNode(q querier, id string) (*models.Node, error) {
	n := models.NewNode()
	var dueDateStr string
	err := q.QueryRow(
		"SELECT id, title, COALESCE(content, ''), COALESCE(due_date, ''), created_at, updated_at FROM nodes WHERE id = ?", id,
	).Scan(&n.ID, &n.Title, &n.Content, &dueDateStr, &n.CreatedAt, &n.UpdatedAt)
	if err != nil {
		return nil, err
	}
	if dueDateStr != "" {
		if d, err := models.ParseDate(dueDateStr); err == nil {
			n.DueDate = &d
		}
	}
	rows, err := q.Query("SELECT rel_name, to_id FROM rels WHERE from_id = ?", id)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var relName, toID string
		rows.Scan(&relName, &toID)
		if toID == "" {
			n.AddTag(relName)
		} else {
			n.AddRelation(models.RelType(relName), toID)
		}
	}
	return n, nil
}

func updateNode(q querier, id, title, content string, dueDate *time.Time, updatedAt string) error {
	var dd any
	if dueDate != nil {
		dd = dueDate.UTC().Format("2006-01-02")
	}
	_, err := q.Exec(
		"UPDATE nodes SET title = ?, content = ?, due_date = ?, updated_at = ? WHERE id = ?",
		title, content, dd, updatedAt, id,
	)
	return err
}

func deleteNode(q querier, id string) error {
	_, err := q.Exec("DELETE FROM nodes WHERE id = ?", id)
	return err
}

func nodeExists(q querier, id string) (bool, error) {
	var e bool
	err := q.QueryRow("SELECT EXISTS(SELECT 1 FROM nodes WHERE id = ?)", id).Scan(&e)
	return e, err
}

func (s *GraphStoreSqlite) AddNode(id, title, content string, dueDate *time.Time, createdAt, updatedAt string) error {
	return addNode(s.db, id, title, content, dueDate, createdAt, updatedAt)
}
func (s *GraphStoreSqlite) GetNode(id string) (*models.Node, error) { return getNode(s.db, id) }
func (s *GraphStoreSqlite) UpdateNode(id, title, content string, dueDate *time.Time, updatedAt string) error {
	return updateNode(s.db, id, title, content, dueDate, updatedAt)
}
func (s *GraphStoreSqlite) DeleteNode(id string) error         { return deleteNode(s.db, id) }
func (s *GraphStoreSqlite) NodeExists(id string) (bool, error) { return nodeExists(s.db, id) }

func (s *txStore) AddNode(id, title, content string, dueDate *time.Time, createdAt, updatedAt string) error {
	return addNode(s.tx, id, title, content, dueDate, createdAt, updatedAt)
}
func (s *txStore) GetNode(id string) (*models.Node, error) { return getNode(s.tx, id) }
func (s *txStore) UpdateNode(id, title, content string, dueDate *time.Time, updatedAt string) error {
	return updateNode(s.tx, id, title, content, dueDate, updatedAt)
}
func (s *txStore) DeleteNode(id string) error         { return deleteNode(s.tx, id) }
func (s *txStore) NodeExists(id string) (bool, error) { return nodeExists(s.db, id) }

// --- ID generation ---

func genID(length int) string {
	b := make([]byte, length)
	for i := range b {
		b[i] = "abcdefghijklmnopqrstuvwxyz"[rand.Intn(26)]
	}
	return string(b)
}

func generateID(q querier) (string, error) {
	for length := 3; ; length++ {
		// Try up to 10*capacity attempts at this length before growing.
		capacity := 1
		for i := 0; i < length; i++ {
			capacity *= 26
		}
		for attempt := 0; attempt < 10*capacity; attempt++ {
			id := genID(length)
			exists, err := nodeExists(q, id)
			if err != nil {
				return "", err
			}
			if !exists {
				return id, nil
			}
		}
	}
}

func (s *GraphStoreSqlite) GenerateID() (string, error) { return generateID(s.db) }
func (s *txStore) GenerateID() (string, error)          { return generateID(s.db) }

// --- Rel operations ---

func addRel(q querier, nodeID, relName, toID string) error {
	_, err := q.Exec("INSERT OR IGNORE INTO rels (from_id, rel_name, to_id) VALUES (?, ?, ?)", nodeID, relName, toID)
	return err
}

func removeRel(q querier, nodeID, relName, toID string) error {
	_, err := q.Exec("DELETE FROM rels WHERE from_id = ? AND rel_name = ? AND to_id = ?", nodeID, relName, toID)
	return err
}

func (s *GraphStoreSqlite) AddRel(nodeID, relName, toID string) error {
	return addRel(s.db, nodeID, relName, toID)
}
func (s *GraphStoreSqlite) RemoveRel(nodeID, relName, toID string) error {
	return removeRel(s.db, nodeID, relName, toID)
}
func (s *txStore) AddRel(nodeID, relName, toID string) error {
	return addRel(s.tx, nodeID, relName, toID)
}
func (s *txStore) RemoveRel(nodeID, relName, toID string) error {
	return removeRel(s.tx, nodeID, relName, toID)
}

// --- FindNodes ---

func findNodes(q querier, filters []*models.Rel) ([]*models.Node, error) {
	query := "SELECT DISTINCT n.id FROM nodes n"
	var joins []string
	var args []any

	for i, f := range filters {
		alias := fmt.Sprintf("r%d", i)
		if f.Target == "" {
			// Tag/property filter: match rels with empty to_id by rel_name prefix.
			joins = append(joins, fmt.Sprintf(
				"JOIN rels %s ON n.id = %s.from_id AND %s.to_id = '' AND %s.rel_name LIKE ? || '%%'",
				alias, alias, alias, alias,
			))
			args = append(args, string(f.Type))
		} else {
			// Edge filter: match rel by exact rel_name and to_id.
			joins = append(joins, fmt.Sprintf(
				"JOIN rels %s ON n.id = %s.from_id AND %s.rel_name = ? AND %s.to_id = ?",
				alias, alias, alias, alias,
			))
			args = append(args, string(f.Type), f.Target)
		}
	}

	if len(joins) > 0 {
		query += " " + strings.Join(joins, " ")
	}
	query += " ORDER BY n.created_at DESC"

	rows, err := q.Query(query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var ids []string
	for rows.Next() {
		var id string
		if err := rows.Scan(&id); err != nil {
			return nil, err
		}
		ids = append(ids, id)
	}

	var nodes []*models.Node
	for _, id := range ids {
		n, err := getNode(q, id)
		if err != nil {
			return nil, err
		}
		nodes = append(nodes, n)
	}
	return nodes, nil
}

func (s *GraphStoreSqlite) FindNodes(filters []*models.Rel) ([]*models.Node, error) {
	return findNodes(s.db, filters)
}

func (s *txStore) FindNodes(filters []*models.Rel) ([]*models.Node, error) {
	return findNodes(s.tx, filters)
}
