package main

import (
	"bytes"
	"context"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/big"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"neuralhunt/internal/core"
)

func pad32Int(x *big.Int) []byte {
	b := x.Bytes()
	out := make([]byte, 32)
	if len(b) > 32 {
		b = b[len(b)-32:]
	}
	copy(out[32-len(b):], b)
	return out
}

func signRaw(key *ecdsa.PrivateKey, message string) (string, error) {
	h := sha256.Sum256([]byte(message))
	r, s, err := ecdsa.Sign(rand.Reader, key, h[:])
	if err != nil {
		return "", err
	}
	raw := append(pad32Int(r), pad32Int(s)...)
	return base64.RawURLEncoding.EncodeToString(raw), nil
}

type apiError struct {
	Status int
	Body   string
}

func (e *apiError) Error() string {
	return fmt.Sprintf("HTTP %d: %s", e.Status, strings.TrimSpace(e.Body))
}

func (e *apiError) Code() string {
	var v struct {
		Code string `json:"code"`
	}
	if json.Unmarshal([]byte(e.Body), &v) == nil {
		return v.Code
	}
	return ""
}

type apiClient struct {
	base string
	hc   *http.Client

	authMu sync.Mutex
	mu     sync.RWMutex
	token  string
	cid    string
	id     identityFile
	key    *ecdsa.PrivateKey
}

func newAPI(base string, id identityFile, key *ecdsa.PrivateKey) *apiClient {
	return &apiClient{
		base: strings.TrimRight(base, "/"),
		hc:   &http.Client{Timeout: 15 * time.Second},
		id:   id,
		key:  key,
	}
}

func (c *apiClient) tokenValue() string {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.token
}

func (c *apiClient) clientID() string {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.cid
}

func (c *apiClient) doRaw(ctx context.Context, method, path string, body, out any) error {
	var rd io.Reader
	if body != nil {
		b, err := json.Marshal(body)
		if err != nil {
			return err
		}
		rd = bytes.NewReader(b)
	}
	req, err := http.NewRequestWithContext(ctx, method, c.base+path, rd)
	if err != nil {
		return err
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	if token := c.tokenValue(); token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	resp, err := c.hc.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	if resp.StatusCode/100 != 2 {
		return &apiError{Status: resp.StatusCode, Body: string(b)}
	}
	if out != nil && len(bytes.TrimSpace(b)) != 0 {
		if err := json.Unmarshal(b, out); err != nil {
			return fmt.Errorf("decode %s: %w", path, err)
		}
	}
	return nil
}

// do retries authenticated API calls once after a fresh cryptographic login.
// This matters for unattended hosted workers: a long network outage or an
// expired JWT must not leave a healthy process permanently stuck on 401.
func (c *apiClient) do(ctx context.Context, method, path string, body, out any) error {
	err := c.doRaw(ctx, method, path, body, out)
	var ae *apiError
	if strings.HasPrefix(path, "/api/auth/") || !errors.As(err, &ae) || ae.Status != http.StatusUnauthorized {
		return err
	}
	if loginErr := c.login(ctx); loginErr != nil {
		return fmt.Errorf("session refresh after HTTP 401: %w", loginErr)
	}
	return c.doRaw(ctx, method, path, body, out)
}

func leadingZeroBitsClient(b []byte) int {
	n := 0
	for _, x := range b {
		if x == 0 {
			n += 8
			continue
		}
		for m := byte(0x80); m != 0 && x&m == 0; m >>= 1 {
			n++
		}
		break
	}
	return n
}

func solveProofClient(challenge, cid string, bits int) string {
	if bits <= 0 {
		return ""
	}
	for i := uint64(0); ; i++ {
		counter := strconv.FormatUint(i, 10)
		h := sha256.Sum256([]byte("nh-pow-v1|" + challenge + "|" + cid + "|" + counter))
		if leadingZeroBitsClient(h[:]) >= bits {
			return counter
		}
	}
}
func (c *apiClient) login(ctx context.Context) error {
	c.authMu.Lock()
	defer c.authMu.Unlock()

	var ch struct {
		ClientID        string `json:"client_id"`
		Challenge       string `json:"challenge"`
		ProofOfWorkBits int    `json:"proof_of_work_bits"`
	}
	if err := c.doRaw(ctx, http.MethodPost, "/api/auth/challenge", map[string]any{"public_jwk": c.id.PublicJWK}, &ch); err != nil {
		return fmt.Errorf("challenge: %w", err)
	}
	sig, err := signRaw(c.key, "login|"+ch.Challenge+"|"+ch.ClientID)
	if err != nil {
		return err
	}
	var lg struct {
		Token    string `json:"token"`
		ClientID string `json:"client_id"`
	}
	pow := solveProofClient(ch.Challenge, ch.ClientID, ch.ProofOfWorkBits)
	if err := c.doRaw(ctx, http.MethodPost, "/api/auth/login", map[string]any{
		"public_jwk":            c.id.PublicJWK,
		"challenge":             ch.Challenge,
		"signature":             sig,
		"proof_of_work_counter": pow,
	}, &lg); err != nil {
		return fmt.Errorf("login: %w", err)
	}
	if old := c.clientID(); old != "" && old != lg.ClientID {
		return fmt.Errorf("server returned a different client identity after re-login: %s != %s", lg.ClientID, old)
	}
	c.mu.Lock()
	c.token, c.cid = lg.Token, lg.ClientID
	c.mu.Unlock()
	return nil
}

type taskCard struct {
	ID                      string  `json:"id"`
	DisplayName             string  `json:"display_name"`
	Description             string  `json:"description"`
	RangeBits               int     `json:"range_bits"`
	Paused                  bool    `json:"paused"`
	Revision                int64   `json:"revision"`
	PointCount              int     `json:"point_count"`
	OwnScore                float64 `json:"own_score"`
	OwnRank                 int64   `json:"own_rank"`
	Selected                bool    `json:"selected"`
	GuessMinIntervalSec     *int    `json:"guess_min_interval_sec,omitempty"`
	ClientSubmitIntervalSec *int    `json:"client_submit_interval_sec,omitempty"`
}

type taskDTO struct {
	ID                      string   `json:"id"`
	PublicSeed              string   `json:"public_seed"`
	RangeBits               int      `json:"range_bits"`
	NextSeq                 int64    `json:"next_seq"`
	ServerMinIntervalSec    int      `json:"server_min_interval_sec"`
	ClientSubmitIntervalSec int      `json:"client_submit_interval_sec"`
	DefaultMaxNodes         int      `json:"default_max_nodes"`
	GuessLotteryWindowSec   int      `json:"guess_lottery_window_sec"`
	GuessLotteryMaxAccepted int      `json:"guess_lottery_max_accepted"`
	BeaconHuntEnabled       int      `json:"beacon_hunt_enabled"`
	BeaconBonusWeight       int      `json:"beacon_bonus_weight"`
	BeaconPaths             []string `json:"beacon_paths"`
	Paused                  bool     `json:"paused"`
	Revision                int64    `json:"revision"`
	DisplayName             string   `json:"display_name"`
	Description             string   `json:"description"`
	ParentTaskID            *string  `json:"parent_task_id"`
}

type point struct {
	ClientID   string  `json:"client_id"`
	Score      float64 `json:"score"`
	X          float64 `json:"x"`
	Y          float64 `json:"y"`
	Z          float64 `json:"z"`
	GuessCount int64   `json:"guess_count"`
	Rank       int64   `json:"rank"`
}

type meDTO struct {
	ClientID string   `json:"client_id"`
	Score    float64  `json:"score"`
	Rank     int64    `json:"rank"`
	Wins     int      `json:"wins"`
	Unlocks  []string `json:"unlocks"`
}

type leader struct {
	ClientID      string   `json:"client_id"`
	Wins          int      `json:"wins"`
	BestScore     float64  `json:"best_score"`
	LiveScore     float64  `json:"live_score"`
	GuessCount    int64    `json:"guess_count"`
	Connected     bool     `json:"connected"`
	Unlocks       []string `json:"unlocks"`
	NFTCount      int      `json:"nft_count"`
	NFTTaskID     *string  `json:"nft_task_id,omitempty"`
	NFTPreviewURI *string  `json:"nft_preview_uri,omitempty"`
}

type publicArtifact struct {
	TaskID         string    `json:"task_id"`
	WinnerClientID string    `json:"winner_client_id"`
	RangeBits      int       `json:"range_bits"`
	CompletedAt    time.Time `json:"completed_at"`
	PreviewURI     string    `json:"preview_uri"`
}

type ownedArtifact struct {
	TaskID      string    `json:"task_id"`
	DisplayName string    `json:"display_name"`
	RangeBits   int       `json:"range_bits"`
	CompletedAt time.Time `json:"completed_at"`
	PreviewURI  string    `json:"preview_uri"`
	DownloadURI string    `json:"download_uri"`
}

func (c *apiClient) tasks(ctx context.Context) ([]taskCard, error) {
	var out []taskCard
	err := c.do(ctx, http.MethodGet, "/api/tasks", nil, &out)
	return out, err
}

func (c *apiClient) selectTask(ctx context.Context, taskID string) (taskDTO, error) {
	var out taskDTO
	err := c.do(ctx, http.MethodPost, "/api/tasks/select", map[string]string{"task_id": taskID}, &out)
	return out, err
}

func (c *apiClient) currentTask(ctx context.Context) (taskDTO, error) {
	var out taskDTO
	err := c.do(ctx, http.MethodGet, "/api/tasks/current", nil, &out)
	return out, err
}

func (c *apiClient) points(ctx context.Context, taskID string, limit int) ([]point, error) {
	var out []point
	err := c.do(ctx, http.MethodGet, "/api/tasks/"+url.PathEscape(taskID)+"/points?limit="+strconv.Itoa(limit), nil, &out)
	return out, err
}

func (c *apiClient) me(ctx context.Context) (meDTO, error) {
	var out meDTO
	err := c.do(ctx, http.MethodGet, "/api/me", nil, &out)
	return out, err
}

type hostedLinkCode struct {
	Code      string    `json:"code"`
	ClientID  string    `json:"client_id"`
	ExpiresAt time.Time `json:"expires_at"`
}

func (c *apiClient) hostedLinkCode(ctx context.Context) (hostedLinkCode, error) {
	var out hostedLinkCode
	err := c.do(ctx, http.MethodPost, "/api/me/customer-link", map[string]any{}, &out)
	return out, err
}

func (c *apiClient) leaderboard(ctx context.Context) ([]leader, error) {
	var out []leader
	err := c.do(ctx, http.MethodGet, "/api/leaderboard", nil, &out)
	return out, err
}

func (c *apiClient) artifacts(ctx context.Context, limit int) ([]publicArtifact, error) {
	var out []publicArtifact
	err := c.do(ctx, http.MethodGet, "/api/public/artifacts?limit="+strconv.Itoa(limit), nil, &out)
	return out, err
}

func (c *apiClient) ownedArtifacts(ctx context.Context, limit int) ([]ownedArtifact, error) {
	var out []ownedArtifact
	err := c.do(ctx, http.MethodGet, "/api/me/artifacts?limit="+strconv.Itoa(limit), nil, &out)
	return out, err
}

func (c *apiClient) guess(ctx context.Context, t taskDTO, seq int64, beaconPath string) (bool, error) {
	guess := expectedGuess(t.ID, t.PublicSeed, c.clientID(), seq, t.RangeBits)
	msg := fmt.Sprintf("guess|%s|%d|%s", t.ID, seq, guess)
	if t.BeaconHuntEnabled == 1 && t.GuessLotteryMaxAccepted > 0 {
		beaconPath = strings.ToUpper(strings.TrimSpace(beaconPath))
		msg += "|" + beaconPath
	}
	sig, err := signRaw(c.key, msg)
	if err != nil {
		return false, err
	}
	var correct bool
	err = c.do(ctx, http.MethodPost, "/api/tasks/"+url.PathEscape(t.ID)+"/guess", map[string]any{"seq": seq, "guess": guess, "signature": sig, "beacon_path": beaconPath}, &correct)
	return correct, err
}

func (c *apiClient) dialWSOnce(ctx context.Context, maxNodes int) (*websocket.Conn, error) {
	u, err := url.Parse(c.base)
	if err != nil {
		return nil, err
	}
	scheme := "ws"
	if u.Scheme == "https" {
		scheme = "wss"
	}
	q := url.Values{}
	q.Set("max_nodes", strconv.Itoa(maxNodes))
	wu := scheme + "://" + u.Host + "/api/ws?" + q.Encode()
	h := http.Header{}
	if token := c.tokenValue(); token != "" {
		h.Set("Authorization", "Bearer "+token)
	}
	conn, resp, err := websocket.DefaultDialer.DialContext(ctx, wu, h)
	if err != nil && resp != nil {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		_ = resp.Body.Close()
		return nil, &apiError{Status: resp.StatusCode, Body: string(b)}
	}
	return conn, err
}

func (c *apiClient) dialWS(ctx context.Context, maxNodes int) (*websocket.Conn, error) {
	conn, err := c.dialWSOnce(ctx, maxNodes)
	var ae *apiError
	if !errors.As(err, &ae) || ae.Status != http.StatusUnauthorized {
		return conn, err
	}
	if loginErr := c.login(ctx); loginErr != nil {
		return nil, fmt.Errorf("websocket session refresh: %w", loginErr)
	}
	return c.dialWSOnce(ctx, maxNodes)
}

func (c *apiClient) downloadPreview(ctx context.Context, taskID, dest string) error {
	path := "/api/public/artifacts/" + url.PathEscape(taskID) + "/preview"
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+path, nil)
	if err != nil {
		return err
	}
	resp, err := c.hc.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return &apiError{Status: resp.StatusCode, Body: string(b)}
	}
	if dir := filepath.Dir(dest); dir != "." {
		if err := os.MkdirAll(dir, 0o750); err != nil {
			return err
		}
	}
	f, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = io.Copy(f, io.LimitReader(resp.Body, 32<<20))
	return err
}

func (c *apiClient) downloadOwnedArtifact(ctx context.Context, taskID, dest string) error {
	path := "/api/me/artifacts/" + url.PathEscape(taskID) + "/download"
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+path, nil)
	if err != nil {
		return err
	}
	if token := c.tokenValue(); token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	resp, err := c.hc.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return &apiError{Status: resp.StatusCode, Body: string(b)}
	}
	if dir := filepath.Dir(dest); dir != "." {
		if err := os.MkdirAll(dir, 0o750); err != nil {
			return err
		}
	}
	f, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = io.Copy(f, io.LimitReader(resp.Body, 64<<20))
	return err
}

func expectedGuess(taskID, seed, clientID string, seq int64, bits int) string {
	// Keep this exactly aligned with internal/core.ExpectedGuess without importing
	// implementation details into the terminal UX layer.
	return core.ExpectedGuess(taskID, seed, clientID, seq, bits)
}

var errSwitching = errors.New("task switch in progress")

// registerHostedWorker binds this worker's freshly authenticated cryptographic
// identity to its Customer Service worker record. The endpoint is reachable
// only on the private Docker network and additionally requires the per-worker
// one-time bearer token injected by the Customer Service.
func (c *apiClient) registerHostedWorker(ctx context.Context) error {
	registerURL := strings.TrimSpace(os.Getenv("NEURALHUNT_WORKER_REGISTER_URL"))
	workerID := strings.TrimSpace(os.Getenv("NEURALHUNT_WORKER_ID"))
	token := strings.TrimSpace(os.Getenv("NEURALHUNT_WORKER_REGISTER_TOKEN"))
	if registerURL == "" && workerID == "" && token == "" {
		return nil
	}
	if registerURL == "" || workerID == "" || token == "" {
		return errors.New("hosted worker registration requires NEURALHUNT_WORKER_REGISTER_URL, NEURALHUNT_WORKER_ID and NEURALHUNT_WORKER_REGISTER_TOKEN")
	}
	body, _ := json.Marshal(map[string]string{"worker_id": workerID, "client_id": c.clientID()})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, registerURL, bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+token)
	hc := &http.Client{Timeout: 10 * time.Second}
	resp, err := hc.Do(req)
	if err != nil {
		return fmt.Errorf("hosted worker register: %w", err)
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
	if resp.StatusCode/100 != 2 {
		return fmt.Errorf("hosted worker register HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(b)))
	}
	return nil
}

// hostedWorkerLeaseLoop is a fail-closed control-plane lease. Billing remains
// authoritative in Customer Service; this loop merely prevents a managed
// worker from running indefinitely if the Customer Service/Docker controller
// disappears. After three consecutive missed 20-second renewals it cancels the
// client context and the managed container exits.
func (c *apiClient) hostedWorkerLeaseLoop(ctx context.Context, cancel context.CancelFunc) {
	leaseURL := strings.TrimSpace(os.Getenv("NEURALHUNT_WORKER_LEASE_URL"))
	workerID := strings.TrimSpace(os.Getenv("NEURALHUNT_WORKER_ID"))
	token := strings.TrimSpace(os.Getenv("NEURALHUNT_WORKER_REGISTER_TOKEN"))
	if leaseURL == "" || workerID == "" || token == "" {
		return
	}
	t := time.NewTicker(20 * time.Second)
	defer t.Stop()
	failures := 0
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			body, _ := json.Marshal(map[string]string{"worker_id": workerID})
			req, err := http.NewRequestWithContext(ctx, http.MethodPost, leaseURL, bytes.NewReader(body))
			if err == nil {
				req.Header.Set("Content-Type", "application/json")
				req.Header.Set("Authorization", "Bearer "+token)
				resp, callErr := (&http.Client{Timeout: 8 * time.Second}).Do(req)
				if callErr == nil {
					_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4096))
					_ = resp.Body.Close()
					if resp.StatusCode/100 == 2 {
						failures = 0
						continue
					}
					// Explicit authentication/revocation responses are authoritative,
					// not transient connectivity failures. Stop immediately; Docker's
					// restart policy cannot bypass registration while the DB lease is
					// revoked.
					if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusConflict {
						cancel()
						return
					}
				}
			}
			failures++
			if failures >= 3 {
				cancel()
				return
			}
		}
	}
}