package db

import (
	"crypto/rand"
	"database/sql"
	"encoding/hex"
	"fmt"
	"sync"
	"time"

	"github.com/redis/go-redis/v9"
)

const sessionTTL = 45 * time.Minute
const maxSessions = 64

type Session struct {
	ID       string
	Conn     ConnReq
	mu       sync.Mutex
	sqlPool  map[string]*sql.DB
	redisCli *redis.Client
	tunnel   *SSHTunnel
	LastUsed time.Time
}

var (
	sessionsMu sync.Mutex
	sessions   = map[string]*Session{}
)

func init() {
	go func() {
		t := time.NewTicker(2 * time.Minute)
		defer t.Stop()
		for range t.C {
			expireSessions()
		}
	}()
}

func newSessionID() string {
	var b [16]byte
	if _, err := rand.Read(b[:]); err != nil {
		return fmt.Sprintf("%d", time.Now().UnixNano())
	}
	return hex.EncodeToString(b[:])
}

func OpenSession(c ConnReq) (string, error) {
	c.SQL = ""
	c.Table = ""
	c.Key = ""
	c.Pattern = ""
	c.SessionID = ""

	dial, tun, err := ApplySSHTunnel(c)
	if err != nil {
		return "", err
	}
	if err := testDial(dial); err != nil {
		if tun != nil {
			tun.Close()
		}
		return "", fmt.Errorf("数据库: %w", err)
	}

	sessionsMu.Lock()
	if len(sessions) >= maxSessions {
		sessionsMu.Unlock()
		if tun != nil {
			tun.Close()
		}
		return "", fmt.Errorf("连接会话已满，请关闭不用的连接后重试")
	}
	id := newSessionID()
	// Persist dialed endpoints (127.0.0.1:local) so later ops reuse the tunnel.
	stored := dial
	stored.SSHEnabled = false
	s := &Session{
		ID:       id,
		Conn:     stored,
		sqlPool:  map[string]*sql.DB{},
		tunnel:   tun,
		LastUsed: time.Now(),
	}
	sessions[id] = s
	sessionsMu.Unlock()
	return id, nil
}

func GetSession(id string) (*Session, error) {
	if id == "" {
		return nil, fmt.Errorf("缺少 sessionId")
	}
	sessionsMu.Lock()
	s, ok := sessions[id]
	sessionsMu.Unlock()
	if !ok {
		return nil, fmt.Errorf("会话已失效，请重新连接")
	}
	s.mu.Lock()
	s.LastUsed = time.Now()
	s.mu.Unlock()
	return s, nil
}

func CloseSession(id string) {
	sessionsMu.Lock()
	s, ok := sessions[id]
	if ok {
		delete(sessions, id)
	}
	sessionsMu.Unlock()
	if ok {
		s.closeAll()
	}
}

func expireSessions() {
	cutoff := time.Now().Add(-sessionTTL)
	sessionsMu.Lock()
	var dead []*Session
	for id, s := range sessions {
		s.mu.Lock()
		stale := s.LastUsed.Before(cutoff)
		s.mu.Unlock()
		if stale {
			dead = append(dead, s)
			delete(sessions, id)
		}
	}
	sessionsMu.Unlock()
	for _, s := range dead {
		s.closeAll()
	}
}

func (s *Session) closeAll() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for k, dbh := range s.sqlPool {
		_ = dbh.Close()
		delete(s.sqlPool, k)
	}
	if s.redisCli != nil {
		_ = s.redisCli.Close()
		s.redisCli = nil
	}
	if s.tunnel != nil {
		s.tunnel.Close()
		s.tunnel = nil
	}
}

func (s *Session) ConnFor(req ConnReq) ConnReq {
	c := s.Conn
	// SQLite 的 database 是文件路径，对象树里的 "main" 不能覆盖它
	if req.Database != "" && driverKey(c.Driver) != "sqlite" {
		c.Database = req.Database
	}
	c.SQL = req.SQL
	c.Table = req.Table
	c.Key = req.Key
	c.Pattern = req.Pattern
	c.Limit = req.Limit
	c.SessionID = s.ID
	return c
}

func ResolveConn(req ConnReq) (ConnReq, error) {
	if req.SessionID == "" {
		return req, nil
	}
	s, err := GetSession(req.SessionID)
	if err != nil {
		return req, err
	}
	return s.ConnFor(req), nil
}