package service

import (
	"context"
	"errors"
	"net"
	"strings"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"

	edgeevents "iop/apps/edge/internal/events"
	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

func TestSubmitRunModelQueueUsesRoutePolicyBeforeProviderInstancePolicy(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "prov-node-route-policy", Client: edgeClient})

	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "prov-node-route-policy",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{
					Name:     "ollama-local",
					Enabled:  true,
					Capacity: 1,
					MaxQueue: 10,
				},
			},
		},
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	bus := edgeevents.NewBus()
	svc := New(reg, bus)
	svc.SetNodeStore(store)

	// run 1: fills capacity (capacity is 1)
	res1, err := svc.SubmitRun(context.Background(), SubmitRunRequest{
		ModelGroupKey: "route-policy-group",
		Adapter:       "ollama-local",
		Background:    true,
		MaxQueue:      1,
	})
	if err != nil {
		t.Fatalf("run1 error: %v", err)
	}
	defer res1.Close()

	// run 2: enters queue in a separate goroutine
	ctx2, cancel2 := context.WithCancel(context.Background())
	defer cancel2()
	errCh2 := make(chan error, 1)
	go func() {
		res2, err := svc.SubmitRun(ctx2, SubmitRunRequest{
			ModelGroupKey: "route-policy-group",
			Adapter:       "ollama-local",
			Background:    true,
			MaxQueue:      1,
		})
		if err == nil {
			res2.Close()
		}
		errCh2 <- err
	}()

	// Wait explicitly until run 2 enters the queue
	waitForQueueLen(t, svc.queue, "route-policy-group", 1)

	// run 3: should be rejected immediately because queue is full (max queue = 1, and run2 is in the queue)
	// We use a short timeout context to prevent blocking indefinitely if something fails.
	ctx3, cancel3 := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel3()

	_, err3 := svc.SubmitRun(ctx3, SubmitRunRequest{
		ModelGroupKey: "route-policy-group",
		Adapter:       "ollama-local",
		Background:    true,
		MaxQueue:      1,
	})
	if err3 == nil {
		t.Fatal("expected third run to fail immediately due to queue full")
	}
	if !strings.Contains(err3.Error(), "queue is full") {
		t.Errorf("expected queue is full error, got: %v", err3)
	}

	// clean up run2 and verify it exits
	cancel2()
	select {
	case err2 := <-errCh2:
		if err2 == nil {
			t.Error("expected run2 to fail with context cancelled, got nil")
		} else if !errors.Is(err2, context.Canceled) {
			t.Errorf("expected run2 to fail with context cancelled, got: %v", err2)
		}
	case <-time.After(1 * time.Second):
		t.Error("timeout waiting for run2 goroutine to exit")
	}
}