AetherForge/server/internal/pool/proxy.go

package pool

import (
	"bufio"
	"crypto/tls"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log"
	"math/big"
	"net"
	"strings"
	"sync"
	"time"

	"crypto-miner-server/internal/models"
)

// Stratum protocol message types
type StratumRequest struct {
	ID     int             `json:"id"`
	Method string          `json:"method"`
	Params json.RawMessage `json:"params"`
}

type StratumResponse struct {
	ID     int             `json:"id"`
	Result json.RawMessage `json:"result"`
	Error  interface{}     `json:"error"`
}

type StratumNotification struct {
	Method string          `json:"method"`
	Params json.RawMessage `json:"params"`
}

// Job represents a mining job from the pool
type Job struct {
	ID            string   `json:"job_id"`
	Height        int64    `json:"height"`
	BlockTemplate string   `json:"blocktemplate"`
	Difficulty    int64    `json:"difficulty"`
	SeedHash      string   `json:"seed_hash"`
	Target        string   `json:"target"`
	Blob          string   `json:"blob"`
	Algo          string   `json:"algo"`
}

// ShareSubmit represents a share submission to the pool
type ShareSubmit struct {
	ID     int    `json:"id"`
	Method string `json:"method"`
	Params []string `json:"params"`
}

// Proxy connects to a Monero mining pool via Stratum protocol
// and acts as a bridge between the pool and our agents
type Proxy struct {
	mu            sync.RWMutex
	config        *Config
	conn          net.Conn
	reader        *bufio.Reader
	connected     bool
	requestID     int
	currentJob    *Job
	jobSubscribed bool
	stopCh        chan struct{}
	wg            sync.WaitGroup

	// Callbacks
	onJob      func(job *Job)
	onShare    func(accepted bool, agentID string, jobID string)
	onError    func(err error)

	// Agent share submissions queue
	shareQueue chan *PendingShare
}

type PendingShare struct {
	AgentID string
	JobID   string
	Nonce   string
	Hash    string
}

type Config struct {
	Host     string
	Port     int
	UseTLS   bool
	Wallet   string
	Password string
}

func NewProxy(cfg *Config) *Proxy {
	return &Proxy{
		config:     cfg,
		stopCh:     make(chan struct{}),
		shareQueue: make(chan *PendingShare, 100),
	}
}

// SetCallbacks sets the callbacks for job updates and share results
func (p *Proxy) SetCallbacks(onJob func(job *Job), onShare func(accepted bool, agentID string, jobID string), onError func(err error)) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.onJob = onJob
	p.onShare = onShare
	p.onError = onError
}

// Start connects to the pool and begins processing
func (p *Proxy) Start() error {
	addr := fmt.Sprintf("%s:%d", p.config.Host, p.config.Port)
	log.Printf("[Pool] Connecting to %s (TLS: %v)...", addr, p.config.UseTLS)

	var conn net.Conn
	var err error

	if p.config.UseTLS {
		netDialer := &net.Dialer{Timeout: 30 * time.Second}
		tlsConn, tlsErr := tls.DialWithDialer(netDialer, "tcp", addr, &tls.Config{})
		conn = tlsConn
		err = tlsErr
	} else {
		dialer := net.Dialer{Timeout: 30 * time.Second}
		conn, err = dialer.Dial("tcp", addr)
	}

	if err != nil {
		return fmt.Errorf("failed to connect to pool: %w", err)
	}

	p.mu.Lock()
	p.conn = conn
	p.reader = bufio.NewReader(conn)
	p.connected = true
	p.mu.Unlock()

	log.Printf("[Pool] Connected to %s", addr)

	// Start reader goroutine
	p.wg.Add(1)
	go p.readLoop()

	// Start share submission goroutine
	p.wg.Add(1)
	go p.shareSubmitLoop()

	// Authenticate with the pool
	if err := p.authenticate(); err != nil {
		return fmt.Errorf("failed to authenticate with pool: %w", err)
	}

	return nil
}

// Stop disconnects from the pool
func (p *Proxy) Stop() {
	close(p.stopCh)
	p.mu.Lock()
	if p.conn != nil {
		p.conn.Close()
		p.connected = false
	}
	p.mu.Unlock()
	p.wg.Wait()
	log.Println("[Pool] Disconnected from pool")
}

// IsConnected returns whether the proxy is connected to the pool
func (p *Proxy) IsConnected() bool {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.connected
}

// GetCurrentJob returns the current mining job
func (p *Proxy) GetCurrentJob() *Job {
	p.mu.RLock()
	defer p.mu.RUnlock()
	if p.currentJob == nil {
		return nil
	}
	jobCopy := *p.currentJob
	return &jobCopy
}

// SubmitShare queues a share for submission to the pool
func (p *Proxy) SubmitShare(agentID, jobID, nonce, hash string) {
	p.shareQueue <- &PendingShare{
		AgentID: agentID,
		JobID:   jobID,
		Nonce:   nonce,
		Hash:    hash,
	}
}

func (p *Proxy) authenticate() error {
	p.requestID++

	// Login request
	loginParams := []interface{}{
		p.config.Wallet,
		p.config.Password,
		"crypto-miner-server/1.0",
	}

	paramsData, _ := json.Marshal(loginParams)
	loginReq := StratumRequest{
		ID:     p.requestID,
		Method: "login",
		Params: paramsData,
	}

	data, _ := json.Marshal(loginReq)
	log.Printf("[Pool] Sending login request...")

	if err := p.writeLine(data); err != nil {
		return fmt.Errorf("failed to send login: %w", err)
	}

	return nil
}

func (p *Proxy) readLoop() {
	defer p.wg.Done()

	for {
		select {
		case <-p.stopCh:
			return
		default:
		}

		p.mu.RLock()
		reader := p.reader
		p.mu.RUnlock()

		if reader == nil {
			time.Sleep(100 * time.Millisecond)
			continue
		}

		line, err := reader.ReadString('\n')
		if err != nil {
			log.Printf("[Pool] Read error: %v", err)
			p.mu.Lock()
			p.connected = false
			p.mu.Unlock()

			if p.onError != nil {
				p.onError(fmt.Errorf("pool connection lost: %w", err))
			}

			// Attempt reconnect after delay
			time.Sleep(10 * time.Second)
			go p.reconnect()
			return
		}

		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}

		p.handleMessage([]byte(line))
	}
}

func (p *Proxy) handleMessage(data []byte) {
	// Try to parse as response first
	var resp StratumResponse
	if err := json.Unmarshal(data, &resp); err == nil && resp.ID > 0 {
		p.handleResponse(resp)
		return
	}

	// Try to parse as notification
	var notif StratumNotification
	if err := json.Unmarshal(data, &notif); err == nil && notif.Method != "" {
		p.handleNotification(notif)
		return
	}

	log.Printf("[Pool] Unhandled message: %s", string(data))
}

func (p *Proxy) handleResponse(resp StratumResponse) {
	log.Printf("[Pool] Response ID=%d: %s", resp.ID, string(resp.Result))

	if resp.ID == 1 {
		// Login response
		var loginResult struct {
			ID     string          `json:"id"`
			Job    json.RawMessage `json:"job"`
			Status string          `json:"status"`
		}
		if err := json.Unmarshal(resp.Result, &loginResult); err != nil {
			log.Printf("[Pool] Failed to parse login result: %v", err)
			return
		}

		log.Printf("[Pool] Login successful! Pool ID: %s, Status: %s", loginResult.ID, loginResult.Status)

		// Parse initial job if provided
		if len(loginResult.Job) > 0 {
			p.parseAndSetJob(loginResult.Job)
		}

		// Subscribe for jobs
		p.subscribe()
	}
}

func (p *Proxy) handleNotification(notif StratumNotification) {
	switch notif.Method {
	case "job":
		log.Printf("[Pool] New job received")
		p.parseAndSetJob(notif.Params)

	case "submit":
		// Share submission result
		var submitResult struct {
			ID     int    `json:"id"`
			Result string `json:"result"`
			Status string `json:"status"`
		}
		if err := json.Unmarshal(notif.Params, &submitResult); err != nil {
			log.Printf("[Pool] Failed to parse submit result: %v", err)
			return
		}
		log.Printf("[Pool] Share submission result: %s", submitResult.Status)

	default:
		log.Printf("[Pool] Unknown notification method: %s", notif.Method)
	}
}

func (p *Proxy) parseAndSetJob(data json.RawMessage) {
	var rawJob struct {
		ID            string `json:"job_id"`
		Height        int64  `json:"height"`
		BlockTemplate string `json:"blocktemplate"`
		Difficulty    int64  `json:"difficulty"`
		SeedHash      string `json:"seed_hash"`
		Target        string `json:"target"`
		Blob          string `json:"blob"`
		Algo          string `json:"algo"`
	}

	// Try different field name variations that pools use
	if err := json.Unmarshal(data, &rawJob); err != nil {
		// Try flat params
		var flatParams []json.RawMessage
		if err2 := json.Unmarshal(data, &flatParams); err2 == nil && len(flatParams) >= 1 {
			json.Unmarshal(flatParams[0], &rawJob)
		} else {
			// Try as array of params
			var params [][]json.RawMessage
			if err3 := json.Unmarshal(data, &params); err3 == nil && len(params) >= 1 && len(params[0]) >= 1 {
				json.Unmarshal(params[0][0], &rawJob)
			} else {
				log.Printf("[Pool] Failed to parse job: %s", string(data))
				return
			}
		}
	}

	job := &Job{
		ID:            rawJob.ID,
		Height:        rawJob.Height,
		BlockTemplate: rawJob.BlockTemplate,
		Difficulty:    rawJob.Difficulty,
		SeedHash:      rawJob.SeedHash,
		Target:        rawJob.Target,
		Blob:          rawJob.Blob,
		Algo:          rawJob.Algo,
	}

	// Calculate target from difficulty if not provided
	if job.Target == "" && job.Difficulty > 0 {
		job.Target = p.difficultyToTarget(job.Difficulty)
	}

	p.mu.Lock()
	p.currentJob = job
	p.mu.Unlock()

	log.Printf("[Pool] New job: ID=%s, Height=%d, Difficulty=%d, Algo=%s",
		job.ID, job.Height, job.Difficulty, job.Algo)

	if p.onJob != nil {
		p.onJob(job)
	}
}

func (p *Proxy) subscribe() {
	p.requestID++
	subParams := []string{}
	paramsData, _ := json.Marshal(subParams)

	subReq := StratumRequest{
		ID:     p.requestID,
		Method: "subscribe",
		Params: paramsData,
	}

	data, _ := json.Marshal(subReq)
	log.Printf("[Pool] Subscribing for jobs...")

	if err := p.writeLine(data); err != nil {
		log.Printf("[Pool] Failed to subscribe: %v", err)
	}
}

func (p *Proxy) shareSubmitLoop() {
	defer p.wg.Done()

	for {
		select {
		case <-p.stopCh:
			return
		case share := <-p.shareQueue:
			p.submitShareToPool(share)
		}
	}
}

func (p *Proxy) submitShareToPool(share *PendingShare) {
	p.mu.RLock()
	connected := p.connected
	p.mu.RUnlock()

	if !connected {
		log.Printf("[Pool] Cannot submit share - not connected to pool")
		return
	}

	p.requestID++

	// Submit share to pool
	submitParams := []string{
		p.config.Wallet,
		share.JobID,
		share.Nonce,
		share.Hash,
	}

	paramsData, _ := json.Marshal(submitParams)
	submitReq := StratumRequest{
		ID:     p.requestID,
		Method: "submit",
		Params: paramsData,
	}

	data, _ := json.Marshal(submitReq)
	log.Printf("[Pool] Submitting share for agent %s (job: %s)...", share.AgentID[:min(8, len(share.AgentID))], share.JobID)

	if err := p.writeLine(data); err != nil {
		log.Printf("[Pool] Failed to submit share: %v", err)
		if p.onShare != nil {
			p.onShare(false, share.AgentID, share.JobID)
		}
		return
	}

	// Read response
	p.mu.RLock()
	reader := p.reader
	p.mu.RUnlock()

	if reader == nil {
		return
	}

	// Note: In a real implementation, we'd read the response asynchronously
	// and match it by ID. For now, we assume accepted.
	if p.onShare != nil {
		p.onShare(true, share.AgentID, share.JobID)
	}
}

func (p *Proxy) reconnect() {
	log.Printf("[Pool] Attempting reconnect in 10 seconds...")
	time.Sleep(10 * time.Second)

	select {
	case <-p.stopCh:
		return
	default:
	}

	if err := p.Start(); err != nil {
		log.Printf("[Pool] Reconnect failed: %v", err)
		if p.onError != nil {
			p.onError(fmt.Errorf("pool reconnect failed: %w", err))
		}
		// Try again
		time.Sleep(30 * time.Second)
		select {
		case <-p.stopCh:
			return
		default:
			go p.reconnect()
		}
	}
}

func (p *Proxy) writeLine(data []byte) error {
	p.mu.RLock()
	conn := p.conn
	p.mu.RUnlock()

	if conn == nil {
		return fmt.Errorf("not connected")
	}

	line := append(data, '\n')
	_, err := conn.Write(line)
	return err
}

func (p *Proxy) difficultyToTarget(difficulty int64) string {
	// Convert difficulty to target hex string
	// target = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF / difficulty
	maxTarget := new(big.Int)
	maxTarget.SetString("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF", 16)

	diff := big.NewInt(difficulty)
	target := new(big.Int).Div(maxTarget, diff)

	// Convert to 32-byte hex (little-endian for Monero)
	bytes := target.Bytes()
	padded := make([]byte, 32)
	copy(padded[32-len(bytes):], bytes)

	// Reverse for little-endian
	for i, j := 0, len(padded)-1; i < j; i, j = i+1, j-1 {
		padded[i], padded[j] = padded[j], padded[i]
	}

	return hex.EncodeToString(padded)
}

// Helper to convert models.Job to pool.Job
func FromModelJob(job *models.Job) *Job {
	if job == nil {
		return nil
	}
	return &Job{
		ID:            job.ID,
		Height:        job.Height,
		BlockTemplate: job.BlockTemplate,
		Difficulty:    job.Difficulty,
		SeedHash:      job.SeedHash,
		Target:        job.Target,
	}
}

// Helper to convert pool.Job to models.Job
func (j *Job) ToModelJob() *models.Job {
	return &models.Job{
		ID:            j.ID,
		Height:        j.Height,
		Difficulty:    j.Difficulty,
		BlockTemplate: j.BlockTemplate,
		SeedHash:      j.SeedHash,
		Target:        j.Target,
		CreatedAt:     time.Now(),
	}
}

// Helper to convert target hex to difficulty
func targetToDifficulty(targetHex string) int64 {
	bytes, err := hex.DecodeString(targetHex)
	if err != nil || len(bytes) == 0 {
		return 0
	}

	// Reverse from little-endian
	for i, j := 0, len(bytes)-1; i < j; i, j = i+1, j-1 {
		bytes[i], bytes[j] = bytes[j], bytes[i]
	}

	target := new(big.Int).SetBytes(bytes)
	if target.Sign() == 0 {
		return 0
	}

	maxTarget := new(big.Int)
	maxTarget.SetString("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF", 16)
	diff := new(big.Int).Div(maxTarget, target)

	return diff.Int64()
}

// ParseBlob extracts fields from a Monero mining blob
func ParseBlob(blobHex string) (map[string]interface{}, error) {
	blob, err := hex.DecodeString(blobHex)
	if err != nil {
		return nil, fmt.Errorf("invalid blob hex: %w", err)
	}

	if len(blob) < 43 {
		return nil, fmt.Errorf("blob too short: %d bytes", len(blob))
	}

	result := make(map[string]interface{})

	// Monero blob structure (simplified):
	// [0:1]   - Reserved (1 byte)
	// [1:9]   - Block ID (8 bytes, little-endian)
	// [9:17]  - Nonce (8 bytes, little-endian) - miners fill this
	// [17:43] - Merkle root + extra data

	result["reserved"] = blob[0]
	result["block_id"] = binary.LittleEndian.Uint64(blob[1:9])
	result["nonce_offset"] = 9
	result["nonce_size"] = 4 // Standard nonce is 4 bytes for most pools

	return result, nil
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}