package repository

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	time "time"

	"github.com/ory/fosite"
	"github.com/sigbit/mcp-auth-proxy/v2/pkg/models"
	"go.etcd.io/bbolt"
)

type kvsRepository struct {
	db         *bbolt.DB
	bucketName string
}

var RefreshTokenGracePeriod = 1 * time.Hour

func NewKVSRepository(path string, bucketName string) (Repository, error) {
	db, err := bbolt.Open(path, 0600, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}
	if err := db.Update(func(tx *bbolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte(bucketName))
		return err
	}); err != nil {
		return nil, fmt.Errorf("failed to create bucket: %w", err)
	}
	return &kvsRepository{
		db:         db,
		bucketName: bucketName,
	}, nil
}

func (r *kvsRepository) create(ctx context.Context, key string, value any) error {
	return r.db.Update(func(tx *bbolt.Tx) error {
		bucket := tx.Bucket([]byte(r.bucketName))
		data, err := json.Marshal(value)
		if err != nil {
			return fmt.Errorf("failed to marshal value: %w", err)
		}
		return bucket.Put([]byte(key), data)
	})
}

func (r *kvsRepository) get(ctx context.Context, key string, value any) error {
	return r.db.View(func(tx *bbolt.Tx) error {
		bucket := tx.Bucket([]byte(r.bucketName))
		data := bucket.Get([]byte(key))
		if data == nil {
			return fosite.ErrNotFound
		}
		return json.Unmarshal(data, value)
	})
}

func (r *kvsRepository) delete(ctx context.Context, key string) error {
	return r.db.Update(func(tx *bbolt.Tx) error {
		bucket := tx.Bucket([]byte(r.bucketName))
		return bucket.Delete([]byte(key))
	})
}

func (r *kvsRepository) update(ctx context.Context, key string, value any) error {
	return r.db.Update(func(tx *bbolt.Tx) error {
		bucket := tx.Bucket([]byte(r.bucketName))
		if bucket.Get([]byte(key)) == nil {
			return fosite.ErrNotFound
		}
		data, err := json.Marshal(value)
		if err != nil {
			return fmt.Errorf("failed to marshal value: %w", err)
		}
		return bucket.Put([]byte(key), data)
	})
}

func (r *kvsRepository) CreateAuthorizeCodeSession(ctx context.Context, code string, fositeReq fosite.Requester) error {
	return r.create(ctx, "authorize_code-"+code, models.FromFositeReq(fositeReq))
}

func (r *kvsRepository) GetAuthorizeCodeSession(ctx context.Context, code string, sess fosite.Session) (fosite.Requester, error) {
	var req models.Request
	if err := r.get(ctx, "authorize_code-"+code, &req); err != nil {
		return nil, err
	}
	fositeReq := req.ToFositeReq()
	if err := restoreSession(fositeReq, req.SessionData, sess); err != nil {
		return nil, err
	}
	return fositeReq, nil
}

func (r *kvsRepository) InvalidateAuthorizeCodeSession(ctx context.Context, code string) error {
	return r.delete(ctx, "authorize_code-"+code)
}

func (r *kvsRepository) CreateAccessTokenSession(ctx context.Context, signature string, fositeReq fosite.Requester) error {
	return r.create(ctx, "access_token-"+signature, models.FromFositeReq(fositeReq))
}

func (r *kvsRepository) GetAccessTokenSession(ctx context.Context, signature string, sess fosite.Session) (fosite.Requester, error) {
	var req models.Request
	if err := r.get(ctx, "access_token-"+signature, &req); err != nil {
		return nil, err
	}
	fositeReq := req.ToFositeReq()
	if err := restoreSession(fositeReq, req.SessionData, sess); err != nil {
		return nil, err
	}
	return fositeReq, nil
}

func (r *kvsRepository) DeleteAccessTokenSession(ctx context.Context, signature string) error {
	return r.delete(ctx, "access_token-"+signature)
}

func (r *kvsRepository) CreateRefreshTokenSession(ctx context.Context, signature string, accessSignature string, req fosite.Requester) (err error) {
	return r.create(ctx, "refresh_token-"+signature, models.FromFositeReq(req))
}

func (r *kvsRepository) GetRefreshTokenSession(ctx context.Context, signature string, sess fosite.Session) (fosite.Requester, error) {
	var req models.Request
	if err := r.get(ctx, "refresh_token-"+signature, &req); err != nil {
		return nil, err
	}
	fositeReq := req.ToFositeReq()
	if err := restoreSession(fositeReq, req.SessionData, sess); err != nil {
		return nil, err
	}
	return fositeReq, nil
}

func (r *kvsRepository) DeleteRefreshTokenSession(ctx context.Context, signature string) error {
	return r.delete(ctx, "refresh_token-"+signature)
}

func (r *kvsRepository) RotateRefreshToken(ctx context.Context, requestID string, signature string) error {
	var req models.Request
	if err := r.get(ctx, "refresh_token-"+signature, &req); err != nil {
		return err
	}
	req.RotatedAt = time.Now()
	return r.update(ctx, "refresh_token-"+signature, req)
}

// RevokeRefreshToken revokes a refresh token as specified in:
// https://tools.ietf.org/html/rfc7009#section-2.1
// If the particular
// token is a refresh token and the authorization server supports the
// revocation of access tokens, then the authorization server SHOULD
// also invalidate all access tokens based on the same authorization
// grant (see Implementation Note).
func (r *kvsRepository) RevokeRefreshToken(ctx context.Context, requestID string) error {
	return r.delete(ctx, "refresh_token-"+requestID)
}

// RevokeAccessToken revokes an access token as specified in:
// https://tools.ietf.org/html/rfc7009#section-2.1
// If the token passed to the request
// is an access token, the server MAY revoke the respective refresh
// token as well.
func (r *kvsRepository) RevokeAccessToken(ctx context.Context, requestID string) error {
	return r.delete(ctx, "access_token-"+requestID)
}

func (r *kvsRepository) RegisterClient(ctx context.Context, fositeClient fosite.Client) error {
	return r.create(ctx, "client-"+fositeClient.GetID(), models.FromFositeClient(fositeClient))
}

// GetClient loads the client by its ID or returns an error
// if the client does not exist or another error occurred.
func (r *kvsRepository) GetClient(ctx context.Context, id string) (fosite.Client, error) {
	var client models.Client
	if err := r.get(ctx, "client-"+id, &client); err != nil {
		return nil, err
	}
	return client.ToFositeClient(), nil
}

// ClientAssertionJWTValid returns an error if the JTI is
// known or the DB check failed and nil if the JTI is not known.
func (r *kvsRepository) ClientAssertionJWTValid(ctx context.Context, jti string) error {
	return errors.New("not implemented")
}

// SetClientAssertionJWT marks a JTI as known for the given
// expiry time. Before inserting the new JTI, it will clean
// up any existing JTIs that have expired as those tokens can
// not be replayed due to the expiry.
func (r *kvsRepository) SetClientAssertionJWT(ctx context.Context, jti string, exp time.Time) error {
	return errors.New("not implemented")
}

func (r *kvsRepository) CreatePKCERequestSession(ctx context.Context, signature string, req fosite.Requester) error {
	return r.create(ctx, "pkce_request-"+signature, models.FromFositeReq(req))
}

func (r *kvsRepository) GetPKCERequestSession(ctx context.Context, signature string, sess fosite.Session) (fosite.Requester, error) {
	var req models.Request
	if err := r.get(ctx, "pkce_request-"+signature, &req); err != nil {
		return nil, err
	}
	fositeReq := req.ToFositeReq()
	if err := restoreSession(fositeReq, req.SessionData, sess); err != nil {
		return nil, err
	}
	return fositeReq, nil
}

func (r *kvsRepository) DeletePKCERequestSession(ctx context.Context, signature string) error {
	return r.delete(ctx, "pkce_request-"+signature)
}

func (r *kvsRepository) CreateAuthorizeRequest(ctx context.Context, fositeAR fosite.AuthorizeRequester) error {
	return r.create(ctx, "authorize_request-"+fositeAR.GetID(), models.FromFositeAuthorizeRequest(fositeAR))
}

func (r *kvsRepository) GetAuthorizeRequest(ctx context.Context, requestID string) (fosite.AuthorizeRequester, error) {
	var ar models.AuthorizeRequest
	if err := r.get(ctx, "authorize_request-"+requestID, &ar); err != nil {
		return nil, err
	}
	return ar.ToFositeAuthorizeRequest(), nil
}

func (r *kvsRepository) DeleteAuthorizeRequest(ctx context.Context, requestID string) error {
	return r.delete(ctx, "authorize_request-"+requestID)
}

func (r *kvsRepository) Close() error {
	return r.db.Close()
}