package deepseek

import (
	"context"
	"deepseek/pow"
	"encoding/base64"
	"encoding/json"
	"log"
)

// CreatePowChallengeReq is the request body used to obtain a PoW challenge.
type CreatePowChallengeReq struct {
	TargetPath string `json:"target_path"`
}

// CreatePowChallengeRes contains a server-issued PoW challenge.
type CreatePowChallengeRes struct {
	Challenge PowChallenge `json:"challenge"`
}

// PowChallenge describes a signed DeepSeek proof-of-work challenge.
type PowChallenge struct {
	Algorithm   string `json:"algorithm"`
	Challenge   string `json:"challenge"`
	Salt        string `json:"salt"`
	Signature   string `json:"signature"`
	Difficulty  uint64 `json:"difficulty"`
	ExpireAt    uint64 `json:"expire_at"`
	ExpireAfter int    `json:"expire_after"`
	TargetPath  string `json:"target_path"`
}

func (c PowChallenge) solve() (uint64, error) {
	answer, err := pow.Challenge{
		Algorithm:  c.Algorithm,
		Challenge:  c.Challenge,
		Salt:       c.Salt,
		Difficulty: c.Difficulty,
		ExpireAt:   c.ExpireAt,
	}.Solve()
	if err != nil {
		return 0, err
	}

	log.Printf("PoW answer: %d", answer)
	return answer, nil
}

// ToBase64 solves c and returns the URL-safe base64-encoded JSON response sent
// in the x-ds-pow-response header.
func (c PowChallenge) ToBase64() ([]byte, error) {
	type s struct {
		Algorithm  string `json:"algorithm"`
		Challenge  string `json:"challenge"`
		Salt       string `json:"salt"`
		Answer     uint64 `json:"answer"`
		Signature  string `json:"signature"`
		TargetPath string `json:"target_path"`
	}

	answer, err := c.solve()
	if err != nil {
		return nil, err
	}
	data, err := json.Marshal(s{
		Algorithm:  c.Algorithm,
		Challenge:  c.Challenge,
		Salt:       c.Salt,
		Answer:     answer,
		Signature:  c.Signature,
		TargetPath: c.TargetPath,
	})
	if err != nil {
		return nil, err
	}

	return []byte(base64.URLEncoding.EncodeToString(data)), nil
}

// CreatePowChallenge requests a PoW challenge using a background context.
func (api *Client) CreatePowChallenge() (CreatePowChallengeRes, error) {
	return api.CreatePowChallengeWithContext(context.Background())
}

// CreatePowChallengeWithContext requests a PoW challenge and honors ctx
// cancellation.
func (api *Client) CreatePowChallengeWithContext(ctx context.Context) (CreatePowChallengeRes, error) {
	var zero CreatePowChallengeRes

	body := CreatePowChallengeReq{TargetPath: "/api/v0/chat/completion"}
	req := NewRequest[CreatePowChallengeRes]("POST", "chat/create_pow_challenge", body)
	resp, err := req.DoWithContext(ctx, api)
	if err != nil {
		return zero, err
	}
	defer resp.Body.Close()
	response, err := req.unmarshallBizResponse(resp)
	if err != nil {
		return zero, err
	}
	return response.Data.Data, nil
}