package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"sync"
	"time"
)

type FatalError struct {
	Reason string
}

func (e *FatalError) Error() string {
	return "the server ended this device for good: " + e.Reason
}

type outcome struct {
	fatal     string
	fast      bool
	stable    bool
	reason    string
	holdUntil time.Time
}

type Session struct {
	cfg     Config
	policy  *Policy
	limiter *Limiter
	fetcher *Fetcher
	stats   *Stats
	log     *Logger
	httpc   *http.Client
	jitter  func() float64
	now     func() time.Time
	sleep   func(ctx context.Context, d time.Duration)
	dialWS  func(ctx context.Context, rawURL string) (*Conn, error)
	limits  ServerLimits
	maxBody int
}

func NewSession(cfg Config, policy *Policy, limiter *Limiter, fetcher *Fetcher, stats *Stats, logger *Logger) *Session {
	return &Session{
		cfg:     cfg,
		policy:  policy,
		limiter: limiter,
		fetcher: fetcher,
		stats:   stats,
		log:     logger,
		httpc:   &http.Client{Timeout: time.Duration(TicketTimeoutSeconds) * time.Second, Transport: &http.Transport{Proxy: http.ProxyFromEnvironment}},
		jitter:  rand.Float64,
		now:     time.Now,
		sleep:   sleepContext,
		dialWS:  Dial,
		limits:  defaultLimits(),
		maxBody: DefaultMaxBodyBytes,
	}
}

func defaultLimits() ServerLimits {
	return ServerLimits{
		MaxBodyBytes:        DefaultMaxBodyBytes,
		PingSeconds:         DefaultPingSeconds,
		ReconnectMinSeconds: DefaultReconnectMin,
		ReconnectMaxSeconds: DefaultReconnectMax,
		JobTimeoutSeconds:   DefaultJobTimeout,
		MaxFrameBytes:       MaxInboundMessage,
	}
}

func (s *Session) Backoff(attempt int) time.Duration {
	low := s.limits.ReconnectMinSeconds
	high := s.limits.ReconnectMaxSeconds
	if low < 1 {
		low = DefaultReconnectMin
	}
	if high < low {
		high = low
	}
	if attempt > 16 {
		attempt = 16
	}
	ceiling := low << attempt
	if ceiling > high || ceiling <= 0 {
		ceiling = high
	}
	seconds := float64(low) + s.jitter()*float64(ceiling-low)
	return time.Duration(seconds * float64(time.Second))
}

func sleepContext(ctx context.Context, d time.Duration) {
	timer := time.NewTimer(d)
	defer timer.Stop()
	select {
	case <-ctx.Done():
	case <-timer.C:
	}
}

func (s *Session) holdDelay(until time.Time) time.Duration {
	wait := until.Sub(s.now())
	wait = max(wait, time.Duration(SuspendedMinSeconds)*time.Second)
	wait = min(wait, time.Duration(SuspendedMaxSeconds)*time.Second)
	return wait + time.Duration(s.jitter()*float64(SuspendedJitterSeconds)*float64(time.Second))
}

func (s *Session) suspension(until string) outcome {
	end := s.now().Add(time.Duration(SuspendedDefaultSeconds) * time.Second)
	if moment, err := time.Parse(time.RFC3339, until); err == nil {
		end = moment
	}
	return outcome{reason: "device is suspended", holdUntil: end}
}

func (s *Session) Run(ctx context.Context) error {
	attempt := 0
	for ctx.Err() == nil {
		result := s.connectOnce(ctx)
		if result.fatal != "" {
			return &FatalError{Reason: result.fatal}
		}
		if ctx.Err() != nil {
			break
		}
		if result.stable {
			attempt = 0
		}
		var delay time.Duration
		if !result.holdUntil.IsZero() {
			delay = s.holdDelay(result.holdUntil)
			attempt = 0
			s.log.Printf("device is suspended until %s, waiting %s", result.holdUntil.UTC().Format(time.RFC3339), delay.Round(time.Second))
			s.sleep(ctx, delay)
			continue
		}
		if result.fast {
			delay = time.Duration((0.5 + s.jitter()) * float64(FastRetryMillis) * float64(time.Millisecond))
		} else {
			delay = s.Backoff(attempt)
			attempt++
		}
		s.log.Printf("disconnected (%s), reconnecting in %s", result.reason, delay.Round(100*time.Millisecond))
		s.sleep(ctx, delay)
	}
	return nil
}

type ticketAnswer struct {
	Ticket string `json:"ticket"`
	WSPath string `json:"ws_path"`
}

type errorAnswer struct {
	Error struct {
		Message string `json:"message"`
		Reason  string `json:"reason"`
		Until   string `json:"until"`
	} `json:"error"`
}

func (s *Session) fetchTicket(ctx context.Context) (ticketAnswer, outcome) {
	var answer ticketAnswer
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, s.cfg.BaseURL()+"/nodes/ticket", nil)
	if err != nil {
		return answer, outcome{fatal: "bad server url"}
	}
	request.Header.Set("Authorization", "Bearer "+s.cfg.Secret)
	request.Header.Set("Accept", "application/json")
	request.Header.Set("User-Agent", "devplace-node/"+Version)
	response, err := s.httpc.Do(request)
	if err != nil {
		return answer, outcome{reason: "ticket request failed: " + stripURL(err).Error()}
	}
	defer response.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(response.Body, 64*1024))
	switch {
	case response.StatusCode == http.StatusOK:
		if json.Unmarshal(body, &answer) != nil || answer.Ticket == "" || !strings.HasPrefix(answer.WSPath, "/") {
			return answer, outcome{reason: "ticket answer is malformed"}
		}
		return answer, outcome{}
	case response.StatusCode == http.StatusUnauthorized:
		return answer, outcome{fatal: "unauthorized, the device secret is not valid (revoked or rotated)"}
	case response.StatusCode == http.StatusForbidden:
		var failure errorAnswer
		_ = json.Unmarshal(body, &failure)
		reason := failure.Error.Reason
		if reason == SuspendedReason {
			return answer, s.suspension(failure.Error.Until)
		}
		if FatalReasons[reason] {
			return answer, outcome{fatal: reason}
		}
		return answer, outcome{reason: "ticket refused: " + reason}
	default:
		return answer, outcome{reason: fmt.Sprintf("ticket request answered http %d", response.StatusCode)}
	}
}

func (s *Session) socketURL(answer ticketAnswer) (string, error) {
	base, err := url.Parse(s.cfg.BaseURL())
	if err != nil {
		return "", err
	}
	scheme := "ws"
	if base.Scheme == "https" {
		scheme = "wss"
	}
	return fmt.Sprintf("%s://%s%s?ticket=%s", scheme, base.Host, answer.WSPath, url.QueryEscape(answer.Ticket)), nil
}

func (s *Session) connectOnce(ctx context.Context) outcome {
	answer, failed := s.fetchTicket(ctx)
	if failed.fatal != "" || failed.reason != "" {
		return failed
	}
	target, err := s.socketURL(answer)
	if err != nil {
		return outcome{fatal: "bad server url"}
	}
	conn, err := s.dialWS(ctx, target)
	if err != nil {
		return outcome{reason: err.Error()}
	}
	defer conn.Close()
	sessionCtx, cancel := context.WithCancel(ctx)
	defer cancel()
	go func() {
		<-sessionCtx.Done()
		conn.Close()
	}()
	started := time.Now()
	result := s.converse(sessionCtx, conn)
	if time.Since(started) >= time.Duration(StableSessionSeconds)*time.Second {
		result.stable = true
	}
	return result
}

func (s *Session) sender(conn *Conn) Sender {
	return func(frameType string, fields map[string]any) error {
		payload := map[string]any{"type": frameType}
		for key, value := range fields {
			payload[key] = value
		}
		data, err := json.Marshal(payload)
		if err != nil {
			return err
		}
		limit := MaxInboundMessage
		if frameType == "job_result" {
			limit = MaxResultMessage
		}
		if len(data) > limit {
			return fmt.Errorf("%s frame of %d bytes exceeds the %d byte cap", frameType, len(data), limit)
		}
		return conn.WriteText(string(data))
	}
}

func (s *Session) helloFields() map[string]any {
	return map[string]any{
		"protocol":       ProtocolVersion,
		"client_version": Version,
		"kind":           DeviceKind,
		"platform":       runtime.GOOS + "/" + runtime.GOARCH,
		"modes":          []string{"fetch", "tunnel"},
		"limits":         map[string]any{"max_concurrent": s.limiter.ConcurrencyCap()},
	}
}

func (s *Session) applyWelcome(frame Frame) {
	limits := defaultLimits()
	if frame.Limits != nil {
		given := *frame.Limits
		if given.MaxBodyBytes > 0 {
			limits.MaxBodyBytes = given.MaxBodyBytes
		}
		if given.PingSeconds > 0 {
			limits.PingSeconds = given.PingSeconds
		}
		if given.ReconnectMinSeconds > 0 {
			limits.ReconnectMinSeconds = given.ReconnectMinSeconds
		}
		if given.ReconnectMaxSeconds > 0 {
			limits.ReconnectMaxSeconds = given.ReconnectMaxSeconds
		}
		limits.MaxConcurrentPerNode = given.MaxConcurrentPerNode
		limits.MaxDonorMBPerDay = given.MaxDonorMBPerDay
		if given.JobTimeoutSeconds > 0 {
			limits.JobTimeoutSeconds = given.JobTimeoutSeconds
		}
	}
	s.limits = limits
	s.maxBody = min(limits.MaxBodyBytes, DefaultMaxBodyBytes)
	if rejected := s.policy.Refine(frame.Destinations); len(rejected) > 0 {
		s.log.Printf("server offered destinations outside the baked-in list, ignored: %s", strings.Join(rejected, ", "))
	}
	s.limiter.ApplyServer(&limits, frame.Settings)
}

func (s *Session) converse(ctx context.Context, conn *Conn) outcome {
	send := s.sender(conn)
	if err := send("hello", s.helloFields()); err != nil {
		return outcome{reason: "hello failed: " + err.Error()}
	}
	_ = conn.SetReadDeadline(time.Now().Add(time.Duration(HelloTimeoutSeconds) * time.Second))
	raw, err := conn.ReadMessage()
	if err != nil {
		return s.closeOutcome(err)
	}
	welcome, err := decodeFrame(raw)
	if err != nil {
		return s.protocolFailure(conn, "welcome is not valid json")
	}
	if welcome.Type == "kick" {
		return s.kickOutcome(welcome.Reason, welcome.Until)
	}
	if welcome.Type != "welcome" {
		return s.protocolFailure(conn, "expected welcome, got "+welcome.Type)
	}
	s.applyWelcome(welcome)
	s.log.Debugf("connected as device %s (%s) with %d allowed destinations", welcome.DeviceUID, welcome.DeviceLabel, len(s.policy.Destinations()))

	var jobs sync.WaitGroup
	tunnels := NewTunnels(s.policy, s.limiter, s.stats, send, s.cfg.MaxTunnels)
	background, stop := context.WithCancel(ctx)
	jobs.Add(1)
	go func() {
		defer jobs.Done()
		s.housekeeping(background, send)
	}()
	defer func() {
		stop()
		tunnels.CloseAll()
		jobs.Wait()
		s.limiter.Flush()
	}()

	patience := time.Duration(s.limits.PingSeconds*MissedPingsBeforeDead) * time.Second
	for ctx.Err() == nil {
		_ = conn.SetReadDeadline(time.Now().Add(patience))
		raw, err := conn.ReadMessage()
		if err != nil {
			return s.closeOutcome(err)
		}
		frame, err := decodeFrame(raw)
		if err != nil {
			return s.protocolFailure(conn, "frame is not valid json")
		}
		switch frame.Type {
		case "ping":
			if err := send("pong", nil); err != nil {
				return outcome{reason: "pong failed: " + err.Error()}
			}
		case "job":
			s.startJob(background, &jobs, send, frame)
		case "tunnel_open":
			tunnels.Open(background, frame)
		case "tunnel_data":
			tunnels.Data(frame)
		case "tunnel_close":
			tunnels.Close(frame.ConnID)
		case "settings_update":
			s.limiter.ApplyServer(nil, frame.Settings)
		case "kick":
			return s.kickOutcome(frame.Reason, frame.Until)
		default:
			return s.protocolFailure(conn, "unexpected frame "+frame.Type)
		}
	}
	return outcome{reason: "stopped"}
}

func (s *Session) protocolFailure(conn *Conn, reason string) outcome {
	_ = conn.WriteClose(CloseProtocol, "protocol")
	return outcome{reason: "protocol error: " + reason}
}

func (s *Session) kickOutcome(reason, until string) outcome {
	if reason == SuspendedReason {
		return s.suspension(until)
	}
	if FatalReasons[reason] {
		return outcome{fatal: reason}
	}
	return outcome{reason: "kicked: " + reason}
}

func (s *Session) closeOutcome(err error) outcome {
	var closeErr *CloseError
	if !errors.As(err, &closeErr) {
		return outcome{reason: "connection lost: " + err.Error()}
	}
	if FatalReasons[closeErr.Reason] {
		return outcome{fatal: closeErr.Reason}
	}
	if closeErr.Code == CloseRetryOwner {
		return outcome{fast: true, reason: "server asked for a retry on the owner worker"}
	}
	return outcome{reason: closeErr.Error()}
}

func (s *Session) housekeeping(ctx context.Context, send Sender) {
	ticker := time.NewTicker(time.Duration(UsageIntervalSeconds) * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if line := s.stats.Roll(time.Now()); line != "" {
				s.log.Printf("%s", line)
			}
			up, down := s.stats.Take()
			if up+down > 0 {
				_ = send("usage", map[string]any{"bytes_up": up, "bytes_down": down})
			}
		}
	}
}

func (s *Session) startJob(ctx context.Context, jobs *sync.WaitGroup, send Sender, frame Frame) {
	if frame.JobID == "" || len(frame.JobID) > MaxIDBytes {
		s.log.Printf("job without a usable id ignored")
		return
	}
	if code := s.limiter.Acquire(); code != "" {
		s.stats.Refused()
		_ = send("job_error", map[string]any{"job_id": frame.JobID, "code": code})
		return
	}
	jobs.Add(1)
	go func() {
		defer jobs.Done()
		defer s.limiter.Release()
		s.runJob(ctx, send, frame)
	}()
}

func (s *Session) runJob(ctx context.Context, send Sender, frame Frame) {
	result, err := s.fetcher.Do(ctx, frame, s.maxBody, s.limiter.Remaining())
	if err != nil {
		code := ErrorCode(err)
		var policyErr *PolicyError
		if errors.As(err, &policyErr) {
			s.stats.Refused()
		} else {
			s.stats.Failed()
		}
		s.log.Debugf("job %s failed: %s", frame.JobID, code)
		_ = send("job_error", map[string]any{"job_id": frame.JobID, "code": code})
		return
	}
	up := int64(len(frame.URL) + 256)
	down := int64(len(result.Body))
	s.limiter.AddBytes(up + down)
	s.stats.Job(up, down)
	fields := map[string]any{
		"job_id":     frame.JobID,
		"status":     result.Status,
		"headers":    result.Headers,
		"body_b64":   base64.StdEncoding.EncodeToString(result.Body),
		"truncated":  result.Truncated,
		"elapsed_ms": int(result.Elapsed.Milliseconds()),
	}
	if err := send("job_result", fields); err != nil {
		s.log.Printf("job %s result could not be sent: %v", frame.JobID, err)
	}
}