package stratum

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"forge-mesh/internal/config"
)

// Proxy forwards stratum miner connections to upstream pools with per-agent worker suffixes.
type Proxy struct {
	cfg       config.StratumConfig
	listeners []net.Listener
	active    int64
	wg        sync.WaitGroup
	cancel    context.CancelFunc
}

// New creates a stratum proxy from server config.
func New(cfg config.StratumConfig) *Proxy {
	return &Proxy{cfg: cfg}
}

// ActiveConnections returns the number of live proxied sessions.
func (p *Proxy) ActiveConnections() int64 {
	return atomic.LoadInt64(&p.active)
}

// Start begins listening without requiring the caller to supply a context.
func (p *Proxy) Start() error {
	ctx, cancel := context.WithCancel(context.Background())
	p.cancel = cancel
	return p.StartContext(ctx)
}

// StartContext listens on XMR and RVN ports and relays to upstream pools.
func (p *Proxy) StartContext(ctx context.Context) error {
	endpoints := []struct {
		listen   string
		upstream string
		label    string
	}{
		{p.cfg.XMRListen, p.cfg.UpstreamXMR, "XMR"},
		{p.cfg.RVNListen, p.cfg.UpstreamRVN, "RVN"},
	}

	for _, ep := range endpoints {
		if ep.listen == "" {
			continue
		}
		ln, err := net.Listen("tcp", ep.listen)
		if err != nil {
			p.Stop()
			return fmt.Errorf("listen %s (%s): %w", ep.listen, ep.label, err)
		}
		p.listeners = append(p.listeners, ln)
		log.Printf("stratum: %s proxy listening on %s -> %s", ep.label, ep.listen, displayUpstream(ep.upstream))

		go p.serveListener(ctx, ln, ep.upstream, ep.label)
	}

	if len(p.listeners) == 0 {
		return fmt.Errorf("no stratum listeners configured")
	}
	return nil
}

func displayUpstream(upstream string) string {
	if upstream == "" {
		return "(local accept)"
	}
	return upstream
}

func (p *Proxy) serveListener(ctx context.Context, ln net.Listener, upstream, label string) {
	for {
		conn, err := ln.Accept()
		if err != nil {
			select {
			case <-ctx.Done():
				return
			default:
				if ne, ok := err.(net.Error); ok && ne.Temporary() {
					time.Sleep(50 * time.Millisecond)
					continue
				}
				return
			}
		}

		p.wg.Add(1)
		go func(c net.Conn) {
			defer p.wg.Done()
			p.handleConn(ctx, c, upstream, label)
		}(conn)
	}
}

func (p *Proxy) handleConn(ctx context.Context, miner net.Conn, upstreamAddr, label string) {
	defer miner.Close()
	atomic.AddInt64(&p.active, 1)
	defer atomic.AddInt64(&p.active, -1)

	agentSuffix := agentSuffixFromAddr(miner.RemoteAddr().String())
	minerReader := bufio.NewReader(miner)
	minerWriter := bufio.NewWriter(miner)

	var upstream net.Conn
	var upstreamReader *bufio.Reader
	var upstreamWriter *bufio.Writer

	openUpstream := func() error {
		if upstreamAddr == "" || upstream != nil {
			return nil
		}
		dialer := net.Dialer{Timeout: 10 * time.Second}
		up, err := dialer.DialContext(ctx, "tcp", upstreamAddr)
		if err != nil {
			return err
		}
		upstream = up
		upstreamReader = bufio.NewReader(upstream)
		upstreamWriter = bufio.NewWriter(upstream)
		log.Printf("stratum: %s relay %s <-> %s worker_suffix=%s", label, miner.RemoteAddr(), upstreamAddr, agentSuffix)
		return nil
	}

	for {
		select {
		case <-ctx.Done():
			if upstream != nil {
				_ = upstream.Close()
			}
			return
		default:
		}

		line, err := minerReader.ReadString('\n')
		if err != nil {
			if upstream != nil {
				_ = upstream.Close()
			}
			return
		}
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}

		var req stratumRequest
		_ = json.Unmarshal([]byte(line), &req)
		method := strings.ToLower(req.Method)

		if upstreamAddr != "" && upstream == nil && isStratumMethod(method) {
			if err := openUpstream(); err != nil {
				log.Printf("stratum: %s upstream dial %s: %v", label, upstreamAddr, err)
				writeLocalError(minerWriter, req.ID, "upstream unavailable")
				continue
			}
		}

		if upstream == nil {
			handleLocal(minerWriter, req)
			continue
		}

		outLine := line
		if method == "mining.authorize" {
			outLine = rewriteAuthorize(line, req, agentSuffix)
		}

		if _, err := upstreamWriter.WriteString(outLine + "\n"); err != nil {
			return
		}
		if err := upstreamWriter.Flush(); err != nil {
			return
		}

		respLine, err := upstreamReader.ReadString('\n')
		if err != nil {
			return
		}
		if _, err := minerWriter.WriteString(respLine); err != nil {
			return
		}
		if err := minerWriter.Flush(); err != nil {
			return
		}
	}
}

type stratumRequest struct {
	ID     json.RawMessage   `json:"id"`
	Method string            `json:"method"`
	Params []json.RawMessage `json:"params"`
}

func isStratumMethod(method string) bool {
	switch method {
	case "mining.subscribe", "mining.authorize", "mining.submit", "mining.extranonce.subscribe":
		return true
	default:
		return false
	}
}

func rewriteAuthorize(line string, req stratumRequest, suffix string) string {
	if len(req.Params) == 0 || suffix == "" {
		return line
	}
	var worker string
	if err := json.Unmarshal(req.Params[0], &worker); err != nil {
		return line
	}
	newWorker := WorkerName(worker, suffix)
	req.Params[0], _ = json.Marshal(newWorker)
	out, err := json.Marshal(req)
	if err != nil {
		return line
	}
	return string(out)
}

// WorkerName appends an agent-specific suffix to the pool worker name.
func WorkerName(baseWorker, agentID string) string {
	if agentID == "" {
		return baseWorker
	}
	if strings.Contains(baseWorker, ".") {
		return baseWorker + "." + agentID
	}
	return baseWorker + "." + agentID
}

func agentSuffixFromAddr(remote string) string {
	host, _, err := net.SplitHostPort(remote)
	if err != nil {
		return strings.ReplaceAll(remote, ":", "_")
	}
	return strings.ReplaceAll(host, ".", "_")
}

func handleLocal(w *bufio.Writer, req stratumRequest) {
	method := strings.ToLower(req.Method)
	switch method {
	case "mining.subscribe":
		_ = writeJSON(w, map[string]interface{}{
			"id":     req.ID,
			"result": []interface{}{[]interface{}{"00000000", "00000000", "00000000"}, "00000001", "00000004"},
			"error":  nil,
		})
	case "mining.authorize", "mining.submit":
		_ = writeJSON(w, map[string]interface{}{
			"id":     req.ID,
			"result": true,
			"error":  nil,
		})
	default:
		_ = writeJSON(w, map[string]interface{}{
			"id":     req.ID,
			"result": true,
			"error":  nil,
		})
	}
	_ = w.Flush()
}

func writeLocalError(w *bufio.Writer, id json.RawMessage, msg string) {
	_ = writeJSON(w, map[string]interface{}{
		"id":     id,
		"result": nil,
		"error":  []interface{}{20, msg, nil},
	})
	_ = w.Flush()
}

func writeJSON(w io.Writer, v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	data = append(data, '\n')
	_, err = w.Write(data)
	return err
}

// Stop closes listeners and waits for active relays to finish.
func (p *Proxy) Stop() {
	if p.cancel != nil {
		p.cancel()
	}
	for _, ln := range p.listeners {
		_ = ln.Close()
	}
	p.wg.Wait()
}

// Close is an alias for Stop for server lifecycle hooks.
func (p *Proxy) Close() { p.Stop() }