package providers

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"
)

type Volcengine struct {
	config Config
	client HTTPClient
	now    func() time.Time
}

func NewVolcengine(c Config, client HTTPClient, now func() time.Time) *Volcengine {
	if client == nil {
		client = http.DefaultClient
	}
	if now == nil {
		now = time.Now
	}
	if c.Region == "" {
		c.Region = "cn-north-1"
	}
	if c.Service == "" {
		c.Service = "cv"
	}
	return &Volcengine{c, client, now}
}
func (v *Volcengine) Submit(ctx context.Context, r Request) (Result, error) {
	p := map[string]any{"req_key": requestModel(r, v.config.Model), "prompt": r.Prompt, "image_urls": r.InputURLs}
	for _, key := range []string{"scale", "width", "height", "min_ratio", "max_ratio", "force_single"} {
		if value, exists := r.Settings[key]; exists && value != nil && value != "" {
			p[key] = value
		}
	}
	return v.call(ctx, "CVSync2AsyncSubmitTask", p)
}
func (v *Volcengine) Query(ctx context.Context, id string) (Result, error) {
	return v.QueryModel(ctx, id, v.config.Model)
}
func (v *Volcengine) QueryModel(ctx context.Context, id, model string) (Result, error) {
	queryOptions, _ := json.Marshal(map[string]any{
		"return_url": true,
		"logo_info":  map[string]any{"add_logo": false, "position": 0, "language": 0, "opacity": 1},
	})
	return v.call(ctx, "CVSync2AsyncGetResult", map[string]any{"req_key": requestModel(Request{Model: model}, v.config.Model), "task_id": id, "req_json": string(queryOptions)})
}
func (v *Volcengine) call(ctx context.Context, action string, payload any) (Result, error) {
	body, _ := json.Marshal(payload)
	endpoint, err := url.Parse(v.config.BaseURL)
	if err != nil {
		return Result{}, errors.New("invalid provider base URL")
	}
	q := endpoint.Query()
	q.Set("Action", action)
	q.Set("Version", "2022-08-31")
	endpoint.RawQuery = canonicalQuery(q)
	date := v.now().UTC()
	xdate := date.Format("20060102T150405Z")
	short := xdate[:8]
	hash := sha(body)
	headers := "content-type:application/json\nhost:" + endpoint.Host + "\nx-content-sha256:" + hash + "\nx-date:" + xdate + "\n"
	signed := "content-type;host;x-content-sha256;x-date"
	canonicalPath := endpoint.EscapedPath()
	if canonicalPath == "" {
		canonicalPath = "/"
	}
	canonical := "POST\n" + canonicalPath + "\n" + endpoint.RawQuery + "\n" + headers + "\n" + signed + "\n" + hash
	scope := short + "/" + v.config.Region + "/" + v.config.Service + "/request"
	stringToSign := "HMAC-SHA256\n" + xdate + "\n" + scope + "\n" + sha([]byte(canonical))
	key := hmacBytes([]byte(v.config.SecretAccessKey), short)
	key = hmacBytes(key, v.config.Region)
	key = hmacBytes(key, v.config.Service)
	key = hmacBytes(key, "request")
	signature := hex.EncodeToString(hmacBytes(key, stringToSign))
	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, endpoint.String(), strings.NewReader(string(body)))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Content-Sha256", hash)
	req.Header.Set("X-Date", xdate)
	req.Header.Set("Authorization", "HMAC-SHA256 Credential="+v.config.AccessKeyID+"/"+scope+", SignedHeaders="+signed+", Signature="+signature)
	resp, err := v.client.Do(req)
	if err != nil {
		return Result{}, &ProviderError{Operation: "volcengine request"}
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
	if resp.StatusCode < 200 || resp.StatusCode >= 300 || !json.Valid(raw) {
		return Result{}, &ProviderError{Operation: "volcengine request", Status: resp.StatusCode}
	}
	r := record(raw)
	d := object(r["data"])
	out := []string{}
	collectURLs(d, &out)
	return Result{TaskID: stringValue(r["task_id"], d["task_id"]), Status: status(first(r["status"], d["status"])), OutputURLs: out, Raw: raw}, nil
}
func sha(b []byte) string { x := sha256.Sum256(b); return hex.EncodeToString(x[:]) }
func hmacBytes(k []byte, s string) []byte {
	h := hmac.New(sha256.New, k)
	_, _ = h.Write([]byte(s))
	return h.Sum(nil)
}
func canonicalQuery(q url.Values) string {
	keys := make([]string, 0, len(q))
	for k := range q {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	parts := []string{}
	for _, k := range keys {
		for _, v := range q[k] {
			parts = append(parts, url.QueryEscape(k)+"="+url.QueryEscape(v))
		}
	}
	return strings.Join(parts, "&")
}