// Package agent implements the Linux client that dials home to the server over
// HTTP long-polling and exposes a root PTY on demand.
package agent

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	lpra "github.com/x-grk/longpolling-remote-access"
	"github.com/x-grk/longpolling-remote-access/internal/proto"
)

// Version is the agent build version (see VERSION.txt).
var Version = lpra.Version

// Config is the agent's runtime configuration (from env / config.env).
type Config struct {
	Host         string // e.g. https://REDACTED:8445
	EnrollSecret string
	Name         string // optional friendly name
	StatePath    string // where the enrolled id+token are persisted
	TunnelAddr   string // host:port for the direct forward tunnel (e.g. panel.example:port); empty disables
	// TunnelPinned keeps TunnelAddr as configured locally (TUNNEL_ADDR /
	// TUNNEL_PORT) instead of following the address set in the panel.
	TunnelPinned bool
}

type state struct {
	AgentID string `json:"agent_id"`
	Token   string `json:"token"`
}

// Agent is a running client.
type Agent struct {
	cfg       Config
	http      *http.Client
	stream    *http.Client
	state     state
	pollCount int

	// Direct tunnel currently running (see setTunnelAddr).
	transport atomic.Value // string: auto, ws or lp, as the server says
	link      atomic.Pointer[wsLink]
	tunMu     sync.Mutex
	tunAddr   string
	tunCancel context.CancelFunc
}

// New builds an agent. It does not perform any network I/O.
func New(cfg Config) *Agent {
	// Tuned transport for low-latency forwarding: keep-alive connection reuse
	// (no TCP/TLS handshake per chunk) and a generous idle pool so many
	// concurrent streams don't fight over 2 connections. Plain HTTP/1.1 keeps
	// each stream on its own connection for predictable parallelism.
	tr := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: (&net.Dialer{
			Timeout:   10 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		ForceAttemptHTTP2:     false,
		TLSNextProto:          map[string]func(string, *tls.Conn) http.RoundTripper{},
		MaxIdleConns:          100,
		MaxIdleConnsPerHost:   32,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: time.Second,
	}
	return &Agent{
		cfg: cfg,
		// Timeout must exceed the server's long-poll hold (25s) plus slack.
		http: &http.Client{Timeout: 60 * time.Second, Transport: tr},
		// File streams take as long as the file needs.
		stream: &http.Client{Transport: tr},
	}
}

// Run enrolls if needed and then loops the control poll forever, reconnecting
// on error with a short backoff. It returns only on a fatal config error.
func (a *Agent) Run(ctx context.Context) error {
	if err := a.loadState(); err != nil {
		return err
	}
	if a.state.AgentID == "" || a.state.Token == "" {
		if err := a.enroll(ctx); err != nil {
			return fmt.Errorf("enroll: %w", err)
		}
	}
	log.Printf("agent %s connected to %s", a.state.AgentID, a.cfg.Host)

	// Direct forward tunnel (low-latency TCP mux + UDP). Its address comes
	// from the panel; fall back to the configured one when the server sends
	// none (tunnel port left at the default on an older server).
	if a.cfg.TunnelPinned || !a.fetchConfig(ctx) {
		a.setTunnelAddr(ctx, a.cfg.TunnelAddr)
	}

	// The control channel: a WebSocket, or long polls. In "auto" a failed or
	// dropped WebSocket gives way to long polling for a minute, then is tried
	// again; "ws" keeps trying the WebSocket; "lp" only polls.
	backoff := time.Second
	for ctx.Err() == nil {
		if a.mode() == proto.TransportLP {
			a.longPoll(ctx, 0)
			continue
		}
		err := a.runWS(ctx)
		if ctx.Err() != nil {
			return nil
		}
		switch a.mode() {
		case proto.TransportWS:
			log.Printf("websocket: %v (retry in %s)", err, backoff)
			select {
			case <-time.After(backoff):
			case <-ctx.Done():
				return nil
			}
			backoff = min(backoff*2, retryCap)
			continue
		case proto.TransportAuto:
			log.Printf("websocket: %v - long polling, websocket again in a minute", err)
		}
		backoff = time.Second
		a.longPoll(ctx, time.Minute)
	}
	return nil
}

// mode is the transport the server set for this agent.
func (a *Agent) mode() string {
	if v, ok := a.transport.Load().(string); ok {
		return v
	}
	return proto.TransportAuto
}

func (a *Agent) applyTransport(v string) {
	switch v {
	case proto.TransportAuto, proto.TransportWS, proto.TransportLP:
		if v != a.mode() {
			log.Printf("transport: %s", v)
			a.transport.Store(v)
		}
	}
}

// longPoll runs control polls: for d, or with d zero until the transport is
// no longer long polling.
func (a *Agent) longPoll(ctx context.Context, d time.Duration) {
	until := time.Now().Add(d)
	backoff := time.Second
	// The server was out of reach (a network down, a machine asleep): once
	// a poll gets through again, "auto" tries the WebSocket at once rather
	// than at the end of the minute - it was the network, not the WebSocket.
	lost := false
	for ctx.Err() == nil {
		if d > 0 && (time.Now().After(until) || a.mode() == proto.TransportWS) {
			return
		}
		if d == 0 && a.mode() != proto.TransportLP {
			return
		}
		// Back from out of reach in "auto": a quick request tells when the
		// server answers again, and the WebSocket is tried then; a poll
		// would sit held on the server first.
		if lost && d > 0 {
			if a.reachable(ctx) {
				return
			}
			select {
			case <-time.After(backoff):
			case <-ctx.Done():
				return
			}
			backoff = min(backoff*2, retryCap)
			continue
		}
		// In "auto" the minute is a minute: a poll still held then is cut
		// short (the server keeps its job for the next one).
		reqCtx, cancel := ctx, context.CancelFunc(func() {})
		if d > 0 {
			reqCtx, cancel = context.WithDeadline(ctx, until)
		}
		err := a.pollOnce(ctx, reqCtx)
		cut := reqCtx.Err() == context.DeadlineExceeded
		cancel()
		if err != nil && cut && ctx.Err() == nil {
			return
		}
		if err != nil {
			log.Printf("poll: %v (retry in %s)", err, backoff)
			select {
			case <-time.After(backoff):
			case <-ctx.Done():
				return
			}
			backoff = min(backoff*2, retryCap)
			lost = true
			continue
		}
		backoff = time.Second
	}
}

// reachable reports whether the server answers a quick request.
func (a *Agent) reachable(ctx context.Context) bool {
	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, a.cfg.Host+"/api/agent/config", nil)
	a.authHeaders(req)
	resp, err := a.http.Do(req)
	if err != nil {
		return false
	}
	resp.Body.Close()
	return resp.StatusCode < 300
}

// retryCap bounds the wait between attempts to reach the server, so a
// machine back from sleep or a network back up is reconnected in seconds.
const retryCap = 5 * time.Second

func (a *Agent) enroll(ctx context.Context) error {
	host, _ := os.Hostname()
	body, _ := json.Marshal(proto.EnrollRequest{
		Name:     a.cfg.Name,
		Hostname: host,
		OS:       runtime.GOOS,
		Arch:     runtime.GOARCH,
		Version:  Version,
		User:     currentUser(),
	})
	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, a.cfg.Host+"/api/agent/enroll", bytes.NewReader(body))
	req.Header.Set("X-Enroll-Secret", a.cfg.EnrollSecret)
	req.Header.Set("Content-Type", "application/json")
	resp, err := a.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
		return fmt.Errorf("server said %s: %s", resp.Status, bytes.TrimSpace(b))
	}
	var er proto.EnrollResponse
	if err := json.NewDecoder(resp.Body).Decode(&er); err != nil {
		return err
	}
	a.state = state{AgentID: er.AgentID, Token: er.Token}
	return a.saveState()
}

// pollOnce performs one control long-poll and dispatches any instruction.
// The poll itself runs on reqCtx, which may end early; what it starts runs on
// ctx.
func (a *Agent) pollOnce(ctx, reqCtx context.Context) error {
	req, _ := http.NewRequestWithContext(reqCtx, http.MethodGet, a.cfg.Host+"/api/agent/poll", nil)
	a.authHeaders(req)
	// Report the version and our OS user occasionally (not every poll), so the
	// panel can flag outdated agents without extra traffic.
	a.pollCount++
	if a.pollCount%20 == 1 {
		req.Header.Set("X-Agent-Version", Version)
		req.Header.Set("X-Agent-User", currentUser())
	}
	resp, err := a.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode == http.StatusUnauthorized {
		return fmt.Errorf("unauthorized (token rejected)")
	}
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("status %s", resp.Status)
	}
	var msg proto.ControlMessage
	if err := json.NewDecoder(resp.Body).Decode(&msg); err != nil {
		return err
	}
	a.applyTunnelHeader(ctx, resp.Header)
	a.applyTransport(resp.Header.Get(proto.TransportHeader))
	a.dispatch(ctx, msg)
	return nil
}

// dispatch starts what a control message asks for, in the background, so the
// control channel keeps listening.
func (a *Agent) dispatch(ctx context.Context, msg proto.ControlMessage) {
	switch msg.Type {
	case proto.CtrlOpenTerminal:
		go a.runTerminal(ctx, msg.SessionID, msg.User)
	case proto.CtrlOpenForward:
		go a.runForward(ctx, msg.SessionID, msg.Port)
	case proto.CtrlTask:
		go a.runTask(ctx, msg.TaskID)
	case proto.CtrlOpenSSH:
		if msg.SSH != nil {
			go a.runSSH(ctx, msg.SessionID, *msg.SSH)
		}
	}
}

// fetchConfig asks the server for the settings it hands to agents and reports
// whether it named a tunnel address.
func (a *Agent) fetchConfig(ctx context.Context) bool {
	ctx2, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	req, _ := http.NewRequestWithContext(ctx2, http.MethodGet, a.cfg.Host+"/api/agent/config", nil)
	a.authHeaders(req)
	resp, err := a.http.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	a.applyTransport(resp.Header.Get(proto.TransportHeader))
	return a.applyTunnelHeader(ctx, resp.Header)
}

// applyTunnelHeader follows the tunnel address the server advertises as
// "host:port"; an empty host means the host we already reach the server on.
func (a *Agent) applyTunnelHeader(ctx context.Context, h http.Header) bool {
	return a.applyTunnelAddr(ctx, h.Get("X-Tunnel-Addr"))
}

func (a *Agent) applyTunnelAddr(ctx context.Context, v string) bool {
	if v == "" || a.cfg.TunnelPinned {
		return false
	}
	host, port, err := net.SplitHostPort(v)
	if err != nil {
		return false
	}
	if host == "" {
		u, err := url.Parse(a.cfg.Host)
		if err != nil || u.Hostname() == "" {
			return false
		}
		host = u.Hostname()
	}
	a.setTunnelAddr(ctx, net.JoinHostPort(host, port))
	return true
}

// setTunnelAddr points the direct tunnel at addr, restarting it if it was
// running against another address. An empty addr stops it.
func (a *Agent) setTunnelAddr(ctx context.Context, addr string) {
	a.tunMu.Lock()
	defer a.tunMu.Unlock()
	if addr == a.tunAddr {
		return
	}
	if a.tunCancel != nil {
		a.tunCancel()
		a.tunCancel = nil
	}
	a.tunAddr = addr
	if addr == "" {
		return
	}
	log.Printf("tunnel: using %s", addr)
	tctx, cancel := context.WithCancel(ctx)
	a.tunCancel = cancel
	go a.runTunnel(tctx, addr)
	go a.runUDP(tctx, addr)
}

func (a *Agent) authHeaders(req *http.Request) {
	req.Header.Set("X-Agent-Id", a.state.AgentID)
	req.Header.Set("Authorization", "Bearer "+a.state.Token)
}

// --- state persistence ---

func (a *Agent) loadState() error {
	b, err := os.ReadFile(a.cfg.StatePath)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return err
	}
	return json.Unmarshal(b, &a.state)
}

func (a *Agent) saveState() error {
	if err := os.MkdirAll(filepath.Dir(a.cfg.StatePath), 0o700); err != nil {
		return err
	}
	b, _ := json.MarshalIndent(a.state, "", "  ")
	return os.WriteFile(a.cfg.StatePath, b, 0o600)
}

var b64 = base64.StdEncoding