// 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"
	"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. REDACTED:8446); 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
	state     state
	pollCount int

	// Direct tunnel currently running (see setTunnelAddr).
	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},
	}
}

// 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)
	}

	backoff := time.Second
	for ctx.Err() == nil {
		if err := a.pollOnce(ctx); err != nil {
			log.Printf("poll: %v (retry in %s)", err, backoff)
			select {
			case <-time.After(backoff):
			case <-ctx.Done():
				return nil
			}
			if backoff < 15*time.Second {
				backoff *= 2
			}
			continue
		}
		backoff = time.Second
	}
	return nil
}

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.
func (a *Agent) pollOnce(ctx context.Context) error {
	req, _ := http.NewRequestWithContext(ctx, 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)
	// Dispatch in the background so we keep polling for more work.
	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)
	}
	return nil
}

// 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()
	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 {
	v := h.Get("X-Tunnel-Addr")
	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