package main

import (
	"encoding/json"
	"errors"
	"net/http"
	"time"

	"github.com/dgraph-io/badger/v4"
)

const cookieDBKey = "cookie:423down"

var errCookieNotFound = errors.New("cookie not found")

type DB struct {
	client *badger.DB
}

func NewDB(path string) (*DB, error) {
	db, err := badger.Open(badger.DefaultOptions(path))
	if err != nil {
		return nil, err
	}
	return &DB{
		client: db,
	}, nil
}

func (d *DB) GetArticleHtml(articleID string) (string, error) {
	var valueCopy []byte
	err := d.client.View(func(txn *badger.Txn) error {
		item, err := txn.Get([]byte("article:" + articleID))
		if err != nil {
			return err
		}

		return item.Value(func(val []byte) error {
			valueCopy = append(valueCopy, val...)
			return nil
		})
	})

	if err != nil {
		return "", err
	}

	return string(valueCopy), nil
}

func (d *DB) SetArticleHtml(articleID string, html string, ttl time.Duration) error {
	if ttl <= 0 {
		return errors.New("ttl 必须大于 0")
	}
	return d.client.Update(func(txn *badger.Txn) error {
		entry := badger.NewEntry([]byte("article:"+articleID), []byte(html)).WithTTL(ttl)
		return txn.SetEntry(entry)
	})
}

// GetCookie 读取缓存中的登录 Cookie。
func (d *DB) GetCookie() ([]*http.Cookie, error) {
	var valueCopy []byte
	err := d.client.View(func(txn *badger.Txn) error {
		item, err := txn.Get([]byte(cookieDBKey))
		if err != nil {
			return err
		}

		return item.Value(func(val []byte) error {
			valueCopy = append(valueCopy, val...)
			return nil
		})
	})
	if errors.Is(err, badger.ErrKeyNotFound) {
		return nil, errCookieNotFound
	}
	if err != nil {
		return nil, err
	}

	var cookies []*http.Cookie
	if err := json.Unmarshal(valueCopy, &cookies); err != nil {
		return nil, err
	}
	if len(cookies) == 0 {
		return nil, errCookieNotFound
	}

	return cookies, nil
}

// SetCookie 保存登录 Cookie，有效期固定为 7 天。
func (d *DB) SetCookie(cookies []*http.Cookie) error {
	if len(cookies) == 0 {
		return errors.New("cookies 不能为空")
	}

	data, err := json.Marshal(cookies)
	if err != nil {
		return err
	}

	return d.client.Update(func(txn *badger.Txn) error {
		entry := badger.NewEntry([]byte(cookieDBKey), data).WithTTL(cookieCacheTTL)
		return txn.SetEntry(entry)
	})
}

// AddCookie 增加或更新单个登录 Cookie，并刷新 7 天有效期。
func (d *DB) AddCookie(cookie *http.Cookie) error {
	if cookie == nil {
		return errors.New("cookie 不能为空")
	}

	cookies, err := d.GetCookie()
	if err != nil && !errors.Is(err, errCookieNotFound) {
		return err
	}

	replaced := false
	for i, existingCookie := range cookies {
		if existingCookie.Name == cookie.Name {
			cookies[i] = cookie
			replaced = true
			break
		}
	}
	if !replaced {
		cookies = append(cookies, cookie)
	}

	return d.SetCookie(cookies)
}

func (d *DB) Close() error {
	return d.client.Close()
}