package grpcserver

import (
	"context"
	"fmt"
	"log"
	"net"
	"time"

	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/checker"
	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/grpc/pb"
	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/models"
	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/services"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/encoding"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/status"
)

func init() {
	encoding.RegisterCodec(JSONCodec{})
}

type vantageServer struct {
	pb.UnimplementedVantageServer
}

// Register carries no licence check, deliberately.
//
// A server row only ever comes from CreateServer, which checks the cap before
// issuing a pre-registration token. By the time an agent calls Register its row
// already exists, so counting here would count the caller itself: an instance
// sitting exactly at its cap would reject the very agent it just authorised, and
// every re-registration after a reinstall would fail too.
//
// The cap is enforced where rows are created, which is the only place it can be
// enforced correctly.
func (s *vantageServer) Register(ctx context.Context, req *pb.RegisterRequest) (*pb.RegisterResponse, error) {
	agentToken, err := services.RegisterServer(req.ServerId, req.PreRegToken, req.Hostname, req.IpAddress, req.OsInfo)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "registration failed: %v", err)
	}
	return &pb.RegisterResponse{AgentToken: agentToken}, nil
}

func (s *vantageServer) SyncKeys(ctx context.Context, req *pb.SyncRequest) (*pb.SyncResponse, error) {
	srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
	if err != nil {
		return nil, status.Errorf(codes.Unauthenticated, "invalid agent token")
	}

	if err := services.UpdateServerLastSeen(srv.ServerID, req.AgentVersion); err != nil {
		log.Printf("failed to update last seen for %s: %v", srv.ServerID, err)
	}

	if err := services.BackfillConsoleConfig(srv); err != nil {
		log.Printf("failed to backfill console config for %s: %v", srv.ServerID, err)
	}

	keys, err := services.BuildAuthorizedKeys(req.ServerId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to build authorized keys: %v", err)
	}

	// Carried on the 30s key poll rather than its own RPC: the agent needs it
	// before its hourly package report, and this is the only message it already
	// receives that often.
	return &pb.SyncResponse{
		PublicKeys:      keys,
		CollectPackages: services.VulnScanningEnabled(srv.InstanceID),
	}, nil
}

func (s *vantageServer) UploadGeneratedKey(ctx context.Context, req *pb.UploadKeyRequest) (*pb.UploadKeyResponse, error) {
	srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
	if err != nil {
		return nil, status.Errorf(codes.Unauthenticated, "invalid agent token")
	}

	key, err := services.CreateKey(srv.InstanceID, req.Label, req.PublicKey, "generated", srv.ServerID, req.PrivateKey, "")
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to store key: %v", err)
	}

	if _, err := services.AssignKey(srv.InstanceID, key.KeyID, srv.ServerID); err != nil {
		log.Printf("failed to auto-assign generated key: %v", err)
	}

	return &pb.UploadKeyResponse{KeyId: key.KeyID}, nil
}

func (s *vantageServer) ReportUpdates(ctx context.Context, req *pb.ReportUpdatesRequest) (*pb.ReportUpdatesResponse, error) {
	srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
	if err != nil {
		return nil, status.Errorf(codes.Unauthenticated, "invalid agent token")
	}

	pkgs := make([]models.PackageUpdate, len(req.Updates))
	for i, u := range req.Updates {
		pkgs[i] = models.PackageUpdate{
			Name:           u.Name,
			CurrentVersion: u.CurrentVersion,
			NewVersion:     u.NewVersion,
		}
	}
	if err := services.StoreAvailableUpdates(srv.ServerID, pkgs); err != nil {
		log.Printf("failed to store updates for %s: %v", srv.ServerID, err)
	}
	return &pb.ReportUpdatesResponse{}, nil
}

// ReportPackages stores a server's installed package set.
//
// It does NOT match against the vulnerability database. Matching happens in
// vulnsched, on the leader: every replica would otherwise need the ~50MB
// database resident, and a database refresh would have N replicas racing to
// rescan the same fleet and sending N digests to the customer.
func (s *vantageServer) ReportPackages(ctx context.Context, req *pb.ReportPackagesRequest) (*pb.ReportPackagesResponse, error) {
	srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
	if err != nil {
		return nil, status.Errorf(codes.Unauthenticated, "invalid agent token")
	}

	// The agent's collect_packages flag is an optimisation; this is the
	// boundary. An agent that ignores the flag still stores nothing.
	if !services.VulnScanningEnabled(srv.InstanceID) {
		return &pb.ReportPackagesResponse{NeedFull: false}, nil
	}

	// The offer call: a hash and no packages. Answering NeedFull=false here is
	// what saves the ~150KB body on the overwhelming majority of reports.
	if len(req.Packages) == 0 {
		known, err := services.HasPackageHash(srv.InstanceID, srv.ServerID, req.Hash)
		if err != nil {
			log.Printf("package hash lookup for %s: %v", srv.ServerID, err)
			return nil, status.Errorf(codes.Internal, "package hash lookup failed")
		}
		return &pb.ReportPackagesResponse{NeedFull: !known}, nil
	}

	pkgs := make([]models.InstalledPackage, len(req.Packages))
	for i, p := range req.Packages {
		pkgs[i] = models.InstalledPackage{
			Name:       p.Name,
			Version:    p.Version,
			Epoch:      int(p.Epoch),
			Arch:       p.Arch,
			SourceName: p.SourceName,
		}
	}

	os := models.OSRelease{
		Family:    req.Os.Family,
		VersionID: req.Os.VersionId,
		Arch:      req.Os.Arch,
	}

	if err := services.StorePackages(srv.InstanceID, srv.ServerID, os, req.Hash, pkgs); err != nil {
		log.Printf("store packages for %s: %v", srv.ServerID, err)
		return nil, status.Errorf(codes.Internal, "failed to store packages")
	}
	return &pb.ReportPackagesResponse{NeedFull: false}, nil
}

// ReportWorkloads stores what a server is running.
//
// It is not gated by licence: the workload registry reads as core fleet
// management rather than a premium add-on. If that ever changes, the check
// belongs here — gating collection, not display — for the same reason it does
// in ReportPackages.
func (s *vantageServer) ReportWorkloads(ctx context.Context, req *pb.ReportWorkloadsRequest) (*pb.ReportWorkloadsResponse, error) {
	srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
	if err != nil {
		return nil, status.Errorf(codes.Unauthenticated, "invalid agent token")
	}

	// The offer call: a hash and no body. Answering NeedFull=false here is what
	// keeps an unchanged 60-second report to one small message.
	//
	// The offer is identified by Full, not by an empty Workloads slice: a host
	// genuinely running nothing sends an empty list as its FULL report, and
	// inferring the offer from emptiness would leave that host answering
	// NeedFull=true forever and never storing anything.
	if !req.Full {
		known, err := services.HasWorkloadHash(srv.InstanceID, srv.ServerID, req.Hash)
		if err != nil {
			log.Printf("workload hash lookup for %s: %v", srv.ServerID, err)
			return nil, status.Errorf(codes.Internal, "workload hash lookup failed")
		}
		return &pb.ReportWorkloadsResponse{NeedFull: !known}, nil
	}

	wls := make([]models.Workload, len(req.Workloads))
	for i, w := range req.Workloads {
		wls[i] = models.Workload{
			Kind:      w.Kind,
			ID:        w.Id,
			Name:      w.Name,
			State:     w.State,
			Health:    w.Health,
			Image:     w.Image,
			Stack:     w.Stack,
			Ports:     w.Ports,
			Restarts:  int(w.Restarts),
			Protected: w.Protected,
		}
		if w.StartedAt != "" {
			if t, err := time.Parse(time.RFC3339, w.StartedAt); err == nil {
				wls[i].StartedAt = t
			}
		}
	}

	if err := storeWorkloadReport(srv.InstanceID, srv.ServerID, req, wls); err != nil {
		log.Printf("store workloads for %s: %v", srv.ServerID, err)
		return nil, status.Errorf(codes.Internal, "failed to store workloads")
	}
	return &pb.ReportWorkloadsResponse{NeedFull: false}, nil
}

func storeWorkloadReport(instanceID, serverID string, req *pb.ReportWorkloadsRequest, wls []models.Workload) error {
	if wls == nil {
		wls = []models.Workload{}
	}
	return services.StoreWorkloads(instanceID, serverID, req.Hash, wls,
		req.DockerOk, req.DockerError, req.SystemdOk, req.SystemdError)
}

func (s *vantageServer) ReportInventory(ctx context.Context, req *pb.InventoryReport) (*pb.InventoryReportResponse, error) {
	srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
	if err != nil {
		return nil, status.Errorf(codes.Unauthenticated, "invalid agent token")
	}
	if err := services.StoreInventory(srv.ServerID, req); err != nil {
		log.Printf("store inventory for %s: %v", srv.ServerID, err)
	}
	return &pb.InventoryReportResponse{}, nil
}

func (s *vantageServer) SyncMonitors(ctx context.Context, req *pb.SyncMonitorsRequest) (*pb.SyncMonitorsResponse, error) {
	srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
	if err != nil {
		return nil, status.Errorf(codes.Unauthenticated, "invalid agent token")
	}
	monitors, err := services.ListMonitorsForRunner(srv.InstanceID, srv.ServerID)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "list monitors")
	}
	specs := make([]pb.MonitorSpec, 0, len(monitors))
	for _, m := range monitors {
		specs = append(specs, pb.MonitorSpec{
			MonitorId:      m.MonitorID,
			Type:           m.Type,
			URL:            m.Target.URL,
			Host:           m.Target.Host,
			Port:           m.Target.Port,
			Method:         m.Target.Method,
			ExpectedStatus: m.Target.ExpectedStatus,
			Keyword:        m.Target.Keyword,
			TLSWarnDays:    m.Target.TLSWarnDays,
			Insecure:       m.Target.Insecure,
			IntervalSec:    m.IntervalSec,
			Retries:        m.Retries,
		})
	}
	return &pb.SyncMonitorsResponse{Monitors: specs}, nil
}

func (s *vantageServer) ReportChecks(ctx context.Context, req *pb.ReportChecksRequest) (*pb.ReportChecksResponse, error) {
	srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
	if err != nil {
		return nil, status.Errorf(codes.Unauthenticated, "invalid agent token")
	}
	for _, r := range req.Results {
		res := checker.Result{Up: r.Up, LatencyMs: r.LatencyMs, Message: r.Message}
		if r.CertExpiryUnix > 0 {
			t := time.Unix(r.CertExpiryUnix, 0)
			res.CertExpiry = &t
		}
		if err := services.IngestResult(srv.InstanceID, srv.ServerID, r.MonitorId, res); err != nil {
			log.Printf("ingest check %s: %v", r.MonitorId, err)
		}
	}
	return &pb.ReportChecksResponse{}, nil
}

func (s *vantageServer) CommandStream(stream pb.Vantage_CommandStreamServer) error {

	msg, err := stream.Recv()
	if err != nil {
		return status.Errorf(codes.InvalidArgument, "expected initial auth message: %v", err)
	}

	srv, err := services.ValidateAgentToken(msg.ServerId, msg.AgentToken)
	if err != nil {
		return status.Errorf(codes.Unauthenticated, "invalid agent token")
	}

	if err := services.UpdateServerLastSeen(srv.ServerID, ""); err != nil {
		log.Printf("update last seen %s: %v", srv.ServerID, err)
	}

	// Serve claims this agent's presence on the bus and subscribes this pod to
	// its command channel, so a dispatch issued by any other replica arrives
	// here. The teardown releases both — an agent that reconnects to a
	// different pod must not leave this one advertising a stream it no longer
	// has.
	ch, release := services.Dispatcher.Serve(stream.Context(), srv.ServerID)
	defer release()

	log.Printf("agent %s connected command stream", srv.ServerID)
	defer log.Printf("agent %s disconnected command stream", srv.ServerID)

	go func() {
		for {
			m, err := stream.Recv()
			if err != nil {
				return
			}
			if m.Result != nil {
				r := m.Result
				log.Printf("agent %s cmd %s: success=%v %s", srv.ServerID, r.CommandId, r.Success, r.Message)
				// Republished so a control action waiting on another pod sees
				// it. Publishing with no subscriber is a no-op, so this is safe
				// for every command result rather than only the awaited ones.
				services.WorkloadResults.DeliverCommand(r)
			}
			if m.WorkloadLogsResult != nil {
				services.WorkloadResults.Deliver(m.WorkloadLogsResult)
			}
			if m.StepResult != nil {
				services.StepResults.Deliver(m.StepResult)
			}
			if m.StepOutput != nil {
				if m.StepOutput.Eof {
					services.StepLogs.Close(m.StepOutput.CommandId)
				} else {
					services.StepLogs.Append(m.StepOutput.CommandId, m.StepOutput.Data)
				}
			}
		}
	}()

	// The heartbeat is what lets the agent tell a live stream from an orphaned
	// one. gRPC keepalive cannot: behind an L7 proxy the agent's connection
	// terminates at the proxy, which answers pings on its own behalf, so a dead
	// pod leaves the agent blocked in Recv forever with commands vanishing into
	// a stream nobody is serving. A message that originates here is the only
	// thing that proves this process is still on the other end.
	ping := time.NewTicker(pingInterval)
	defer ping.Stop()

	// Beats are counted and reported periodically rather than logged one by
	// one: at one every 20s per agent, a fleet of any size would drown every
	// other line in the log. What is worth a line of its own is the first beat
	// (it tells the operator this stream's watchdog is now armed on the agent
	// side) and any failure to send one.
	var beats int
	summary := time.NewTicker(pingSummaryInterval)
	defer summary.Stop()

	ctx := stream.Context()
	for {
		select {
		case <-ctx.Done():
			return nil
		case <-summary.C:
			log.Printf("agent %s command stream healthy, %d beats in the last %s",
				srv.ServerID, beats, pingSummaryInterval)
			beats = 0
		case <-ping.C:
			// A failed send is the point: it is how this side learns the stream
			// is gone, which runs the deferred release and frees the agent's
			// presence claim for whichever pod it reconnects to.
			if err := stream.Send(&pb.ServerCommand{Ping: &pb.PingCmd{}}); err != nil {
				log.Printf("agent %s command stream beat failed after %d beats: %v",
					srv.ServerID, beats, err)
				return err
			}
			beats++
			if beats == 1 {
				log.Printf("agent %s command stream beating every %s", srv.ServerID, pingInterval)
			}
		case cmd, ok := <-ch:
			if !ok {
				return nil
			}
			if err := stream.Send(cmd); err != nil {
				return err
			}
		}
	}
}

// How often the server beats on an idle command stream. Comfortably under the
// agent's staleness threshold, so a single dropped beat does not cost a
// reconnect.
const pingInterval = 20 * time.Second

// How often an otherwise silent healthy stream says so. Long enough that a
// large fleet does not fill the log, short enough that "this pod is still
// serving that agent" is answerable from the log rather than by inference.
const pingSummaryInterval = 5 * time.Minute

// StartGRPC serves the agent API until stop is called.
//
// It returns a stop function rather than serving forever because an abrupt exit
// is not a neutral act here: every CommandStream handler holds an agent's
// presence claim, released by a deferred call that a killed process never runs.
// The claim then outlives its owner for the remainder of its 30s TTL, during
// which dispatch believes the agent is reachable, publishes to a channel with
// no subscriber, and fails as "agent offline" — a pod that has already exited
// still answering for an agent it can no longer reach.
func StartGRPC(port int) (stop func(), err error) {
	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return nil, fmt.Errorf("failed to listen: %w", err)
	}

	s := grpc.NewServer(

		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
			MinTime:             20 * time.Second,
			PermitWithoutStream: false,
		}),
		grpc.KeepaliveParams(keepalive.ServerParameters{

			Time:    45 * time.Second,
			Timeout: 10 * time.Second,
		}),
	)
	pb.RegisterVantageServer(s, &vantageServer{})

	go func() {
		log.Printf("gRPC server listening on :%d", port)
		if err := s.Serve(lis); err != nil {
			log.Fatalf("gRPC server error: %v", err)
		}
	}()

	// GracefulStop sends GOAWAY and waits for the handlers to return, which is
	// what runs those deferred releases and, on the agent's side, ends the
	// stream with a clean error it reconnects from immediately rather than
	// waiting out a TCP timeout.
	//
	// It is bounded: an idle CommandStream returns as soon as its context is
	// cancelled, but a console relay mid-transfer would otherwise hold the
	// process past the pod's grace period and earn a SIGKILL — which is the
	// abrupt exit this exists to avoid.
	return func() {
		done := make(chan struct{})
		go func() {
			s.GracefulStop()
			close(done)
		}()
		select {
		case <-done:
			log.Println("gRPC server stopped gracefully")
		case <-time.After(grpcStopTimeout):
			log.Printf("gRPC server did not stop within %s, forcing", grpcStopTimeout)
			s.Stop()
		}
	}, nil
}

// How long GracefulStop is given before outstanding streams are cut. Comfortably
// inside the chart's termination grace period, so the forced stop below still
// leaves time for the HTTP server to drain.
const grpcStopTimeout = 10 * time.Second
