package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
)

type VerifyResponse struct {
	Sol       string `json:"sol"`
	Evm       string `json:"evm"`
	IsNewUser bool   `json:"isNewUser"`
}

func getAxiomWallet(privateKeyB58 string) (*VerifyResponse, error) {
	kp, err := keypairFromSecretKey(privateKeyB58)
	if err != nil {
		return nil, fmt.Errorf("failed to create keypair: %w", err)
	}

	addr := kp.PublicKeyBase58()

	// Get nonce
	transport, err := getNextProxy()
	if err != nil {
		return nil, fmt.Errorf("failed to get proxy: %w", err)
	}

	client := &http.Client{Transport: transport}

	nonceReq := map[string]string{"walletAddress": addr}
	nonceBody, _ := json.Marshal(nonceReq)

	req, err := http.NewRequest("POST", "https://api.axiom.trade/wallet-nonce", bytes.NewReader(nonceBody))
	if err != nil {
		return nil, fmt.Errorf("failed to create nonce request: %w", err)
	}

	req.Header.Set("accept", "application/json, text/plain, */*")
	req.Header.Set("content-type", "application/json")
	req.Header.Set("origin", "https://axiom.trade")
	req.Header.Set("referer", "https://axiom.trade/")
	req.Header.Set("user-agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36")

	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to get nonce: %w", err)
	}
	defer resp.Body.Close()

	nonceText, _ := io.ReadAll(resp.Body)
	nonce := string(nonceText)

	// Try to parse as JSON
	var nonceData map[string]interface{}
	if json.Unmarshal(nonceText, &nonceData) == nil {
		if n, ok := nonceData["nonce"].(string); ok {
			nonce = n
		}
	}

	// Sign message
	message := fmt.Sprintf("By signing, you agree to Axiom's Terms of Use & Privacy Policy (axiom.trade/legal).\n\nNonce: %s", nonce)
	signature := signMessage(kp.PrivateKey, []byte(message))

	// Verify wallet
	verifyReq := map[string]interface{}{
		"walletAddress":     addr,
		"signature":         signature,
		"nonce":             nonce,
		"referrer":          nil,
		"allowRegistration": true,
		"isVerify":          false,
		"forAddCredential":  false,
		"allowLinking":      false,
	}
	verifyBody, _ := json.Marshal(verifyReq)

	req2, err := http.NewRequest("POST", "https://api.axiom.trade/verify-wallet-v2", bytes.NewReader(verifyBody))
	if err != nil {
		return nil, fmt.Errorf("failed to create verify request: %w", err)
	}

	req2.Header.Set("accept", "application/json, text/plain, */*")
	req2.Header.Set("content-type", "application/json")
	req2.Header.Set("origin", "https://axiom.trade")
	req2.Header.Set("referer", "https://axiom.trade/")
	req2.Header.Set("user-agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36")

	resp2, err := client.Do(req2)
	if err != nil {
		return nil, fmt.Errorf("failed to verify wallet: %w", err)
	}
	defer resp2.Body.Close()

	responseText, _ := io.ReadAll(resp2.Body)

	if resp2.StatusCode != http.StatusOK {
		errMsg := string(responseText)
		if len(errMsg) > 200 {
			errMsg = errMsg[:200]
		}
		return nil, fmt.Errorf("axiom error (%d): %s", resp2.StatusCode, errMsg)
	}

	var result VerifyResponse
	if err := json.Unmarshal(responseText, &result); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	return &result, nil
}

func getAxiomBalance(pk string) (string, float64, error) {
	ax, err := getAxiomWallet(pk)
	if err != nil {
		return "", 0, err
	}

	if ax.IsNewUser {
		return ax.Sol, 0, nil
	}

	cfg := &Config{JupiterRateLimitMs: 1100}
	portfolio, err := getPortfolioValue(ax.Sol, cfg)
	if err != nil {
		return ax.Sol, 0, err
	}

	return ax.Sol, portfolio.Value, nil
}