package service

import (
	"context"
	"net"
	"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/packages/go/execution"
	iop "iop/proto/gen/iop"
)

func TestProviderCandidateResponseStallTimeout(t *testing.T) {
	for _, tc := range []struct {
		name string
		raw  int64
		want int64
	}{
		{name: "omitted defaults", want: execution.DefaultResponseStallTimeoutMS},
		{name: "configured value", raw: 45000, want: 45000},
	} {
		t.Run(tc.name, func(t *testing.T) {
			candidate := candidateNode{}
			applyProviderDispatchFields(&candidate, config.NodeProviderConf{ResponseStallTimeoutMS: tc.raw})
			if got := candidate.responseStallTimeoutMS; got != tc.want {
				t.Errorf("response stall timeout = %d, want %d", got, tc.want)
			}
		})
	}
}

func TestDirectDispatchUsesZeroWireStallTimeout(t *testing.T) {
	t.Run("normalized", func(t *testing.T) {
		edgeConn, nodeConn := net.Pipe()
		t.Cleanup(func() { _ = edgeConn.Close(); _ = nodeConn.Close() })
		parser := toki.ParserMap{toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) { m := &iop.RunRequest{}; return m, proto.Unmarshal(b, m) }}
		edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parser)
		nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parser)
		wires := make(chan *iop.RunRequest, 1)
		toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) { wires <- proto.Clone(req).(*iop.RunRequest) })
		svc := directStallTimeoutService(edgeClient)
		result, err := svc.SubmitRun(context.Background(), SubmitRunRequest{NodeRef: "direct-node", RunID: "direct-run", Adapter: "adapter", Target: "target", Background: true, ResponseStallTimeoutMS: 45000})
		if err != nil {
			t.Fatal(err)
		}
		if got := result.Dispatch().ResponseStallTimeoutMS; got != execution.DefaultResponseStallTimeoutMS {
			t.Fatalf("dispatch timeout = %d", got)
		}
		select {
		case wire := <-wires:
			if got := wire.GetResponseStallTimeoutMs(); got != 0 {
				t.Fatalf("wire timeout = %d, want 0", got)
			}
		case <-time.After(time.Second):
			t.Fatal("did not receive RunRequest")
		}
	})

	t.Run("tunnel", func(t *testing.T) {
		edgeConn, nodeConn := net.Pipe()
		t.Cleanup(func() { _ = edgeConn.Close(); _ = nodeConn.Close() })
		parser := toki.ParserMap{toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.ProviderTunnelRequest{}
			return m, proto.Unmarshal(b, m)
		}}
		edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parser)
		nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parser)
		wires := make(chan *iop.ProviderTunnelRequest, 1)
		toki.AddListenerTyped[*iop.ProviderTunnelRequest](&nodeClient.Communicator, func(req *iop.ProviderTunnelRequest) { wires <- proto.Clone(req).(*iop.ProviderTunnelRequest) })
		svc := directStallTimeoutService(edgeClient)
		result, err := svc.SubmitProviderTunnel(context.Background(), SubmitProviderTunnelRequest{NodeRef: "direct-node", RunID: "direct-tunnel", Adapter: "adapter", Target: "target", ResponseStallTimeoutMS: 45000})
		if err != nil {
			t.Fatal(err)
		}
		defer result.Close()
		if got := result.Dispatch().ResponseStallTimeoutMS; got != execution.DefaultResponseStallTimeoutMS {
			t.Fatalf("dispatch timeout = %d", got)
		}
		select {
		case wire := <-wires:
			if got := wire.GetResponseStallTimeoutMs(); got != 0 {
				t.Fatalf("wire timeout = %d, want 0", got)
			}
		case <-time.After(time.Second):
			t.Fatal("did not receive ProviderTunnelRequest")
		}
	})
}

func directStallTimeoutService(client *toki.TcpClient) *Service {
	registry := edgenode.NewRegistry()
	registry.Register(&edgenode.NodeEntry{NodeID: "direct-node", Client: client, DispatchReady: true})
	return New(registry, edgeevents.NewBus())
}

type timeoutMatrixTestCase struct {
	name            string
	isTunnel        bool
	isQueued        bool
	wantProvID      string
	wantTarget      string
	wantTimeout     int64
	wantExecPath    string
	wantQueueReason string
}

func TestProviderPoolResponseStallTimeoutIdentityMatrix(t *testing.T) {
	tests := []timeoutMatrixTestCase{
		{name: "normalized_immediate", isTunnel: false, isQueued: false, wantProvID: "prov-1", wantTarget: "target-1", wantTimeout: 30000, wantExecPath: "normalized", wantQueueReason: "dispatched"},
		{name: "normalized_queued", isTunnel: false, isQueued: true, wantProvID: "prov-2", wantTarget: "target-2", wantTimeout: 60000, wantExecPath: "normalized", wantQueueReason: "capacity_full"},
		{name: "tunnel_immediate", isTunnel: true, isQueued: false, wantProvID: "prov-1", wantTarget: "target-1", wantTimeout: 30000, wantExecPath: "provider_tunnel", wantQueueReason: "dispatched"},
		{name: "tunnel_queued", isTunnel: true, isQueued: true, wantProvID: "prov-2", wantTarget: "target-2", wantTimeout: 60000, wantExecPath: "provider_tunnel", wantQueueReason: "capacity_full"},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			runTimeoutMatrixSubtest(t, tc)
		})
	}
}

func runTimeoutMatrixSubtest(t *testing.T, tc timeoutMatrixTestCase) {
	edgeConn, nodeConn := net.Pipe()
	t.Cleanup(func() { _ = edgeConn.Close(); _ = nodeConn.Close() })

	provType := "ollama"
	if tc.isTunnel {
		provType = "vllm"
	}

	runWires := make(chan *iop.RunRequest, 2)
	tunnelWires := make(chan *iop.ProviderTunnelRequest, 2)
	edgeClient, _ := setupTimeoutMatrixClients(edgeConn, nodeConn, tc.isTunnel, runWires, tunnelWires)

	groupKey := "group-timeout-identity"
	svc, store, catalog, policy := setupTimeoutMatrixService(edgeClient, provType, groupKey)

	resDispatch, closeResult := executeTimeoutMatrixSubmit(t, svc, store, catalog, policy, groupKey, provType, tc)
	if closeResult != nil {
		defer closeResult()
	}

	assertTimeoutMatrixDispatch(t, resDispatch, tc)
	assertTimeoutMatrixWire(t, tc, runWires, tunnelWires)

	if closeResult != nil {
		closeResult()
		closeResult = nil
	}
	svc.HandleNodeDisconnect("node-timeout-matrix", 0, "test-cleanup")
	assertQueueSettled(t, svc.queue)
}

func setupTimeoutMatrixClients(edgeConn, nodeConn net.Conn, isTunnel bool, runWires chan *iop.RunRequest, tunnelWires chan *iop.ProviderTunnelRequest) (*toki.TcpClient, *toki.TcpClient) {
	var parser toki.ParserMap
	if isTunnel {
		parser = toki.ParserMap{
			toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
				m := &iop.ProviderTunnelRequest{}
				return m, proto.Unmarshal(b, m)
			},
		}
	} else {
		parser = toki.ParserMap{
			toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
				m := &iop.RunRequest{}
				return m, proto.Unmarshal(b, m)
			},
		}
	}

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parser)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parser)

	if isTunnel {
		toki.AddListenerTyped[*iop.ProviderTunnelRequest](&nodeClient.Communicator, func(req *iop.ProviderTunnelRequest) {
			tunnelWires <- proto.Clone(req).(*iop.ProviderTunnelRequest)
		})
	} else {
		toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) {
			runWires <- proto.Clone(req).(*iop.RunRequest)
		})
	}
	return edgeClient, nodeClient
}

func buildTimeoutMatrixStore(provType, health1 string) *edgenode.NodeStore {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-timeout-matrix",
		Runtime: config.RuntimeConf{Concurrency: 2},
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{{Name: "shared-adapter", Enabled: true}},
			VllmInstances:   []config.VllmInstanceConf{{Name: "shared-adapter", Enabled: true}},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-1", Type: provType, Adapter: "shared-adapter", Models: []string{"target-1"}, Health: health1, Capacity: 1, ResponseStallTimeoutMS: 30000},
			{ID: "prov-2", Type: provType, Adapter: "shared-adapter", Models: []string{"target-2"}, Health: "available", Capacity: 1, ResponseStallTimeoutMS: 60000},
		},
	})
	return store
}

func setupTimeoutMatrixService(edgeClient *toki.TcpClient, provType, groupKey string) (*Service, *edgenode.NodeStore, []config.ModelCatalogEntry, groupPolicy) {
	catalog := []config.ModelCatalogEntry{
		{ID: groupKey, Providers: map[string]string{"prov-1": "target-1", "prov-2": "target-2"}},
	}
	store := buildTimeoutMatrixStore(provType, "available")
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:         "node-timeout-matrix",
		LifecycleState: edgenode.LifecycleConnected,
		Client:         edgeClient,
		DispatchReady:  true,
	})
	svc := New(reg, edgeevents.NewBus())
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)
	policy := groupPolicyFromStore(store, reg.AllReady(), "shared-adapter", "target-1")
	return svc, store, catalog, policy
}

func executeTimeoutMatrixSubmit(t *testing.T, svc *Service, store *edgenode.NodeStore, catalog []config.ModelCatalogEntry, policy groupPolicy, groupKey, provType string, tc timeoutMatrixTestCase) (RunDispatch, func()) {
	runID := "run-" + tc.name
	if !tc.isQueued {
		if tc.isTunnel {
			res, err := svc.SubmitProviderTunnel(context.Background(), SubmitProviderTunnelRequest{RunID: runID, ModelGroupKey: groupKey, ProviderPool: true})
			if err != nil {
				t.Fatalf("immediate tunnel submit error: %v", err)
			}
			return res.Dispatch(), res.Close
		}
		res, err := svc.SubmitRun(context.Background(), SubmitRunRequest{RunID: runID, ModelGroupKey: groupKey, ProviderPool: true, Background: true})
		if err != nil {
			t.Fatalf("immediate normalized submit error: %v", err)
		}
		return res.Dispatch(), res.Close
	}

	cands, pol, err := svc.resolveProviderPoolCandidates(SubmitRunRequest{ModelGroupKey: groupKey, ProviderPool: true}, store, catalog)
	if err != nil || len(cands) < 2 {
		t.Fatalf("resolve candidates: err=%v len=%d", err, len(cands))
	}
	sel1, _, err1 := svc.queue.admitWithReason(t.Context(), groupKey, "shared-adapter", "target-1", cands, pol, nil, false, true)
	if err1 != nil {
		t.Fatalf("admit prov-1: %v", err1)
	}
	r1 := newQueueReservation(svc.queue, sel1)

	sel2, _, err2 := svc.queue.admitWithReason(t.Context(), groupKey, "shared-adapter", "target-2", cands, pol, nil, false, true)
	if err2 != nil {
		r1.release("cleanup-prov1")
		t.Fatalf("admit prov-2: %v", err2)
	}
	r2 := newQueueReservation(svc.queue, sel2)

	type submitOut struct {
		dispatch RunDispatch
		close    func()
		err      error
	}
	outCh := make(chan submitOut, 1)

	go func() {
		if tc.isTunnel {
			res, err := svc.SubmitProviderTunnel(context.Background(), SubmitProviderTunnelRequest{RunID: runID, ModelGroupKey: groupKey, ProviderPool: true})
			if err != nil {
				outCh <- submitOut{err: err}
				return
			}
			outCh <- submitOut{dispatch: res.Dispatch(), close: res.Close}
		} else {
			res, err := svc.SubmitRun(context.Background(), SubmitRunRequest{RunID: runID, ModelGroupKey: groupKey, ProviderPool: true, Background: true})
			if err != nil {
				outCh <- submitOut{err: err}
				return
			}
			outCh <- submitOut{dispatch: res.Dispatch(), close: res.Close}
		}
	}()

	requireProviderPoolPending(t, svc.queue, 1)
	store2 := buildTimeoutMatrixStore(provType, "disabled")
	svc.SetRuntimeConfig(store2, catalog, policy)
	requireProviderPoolPending(t, svc.queue, 1)
	r2.release("make-prov2-available")

	select {
	case out := <-outCh:
		r1.release("cleanup-prov1")
		if out.err != nil {
			t.Fatalf("queued submit error: %v", out.err)
		}
		return out.dispatch, out.close
	case <-time.After(3 * time.Second):
		r1.release("cleanup-prov1")
		t.Fatal("timed out waiting for queued submit result")
		return RunDispatch{}, nil
	}
}

func assertTimeoutMatrixDispatch(t *testing.T, disp RunDispatch, tc timeoutMatrixTestCase) {
	runID := "run-" + tc.name
	if got := disp.RunID; got != runID {
		t.Errorf("RunID = %q, want %q", got, runID)
	}
	if got := disp.ProviderID; got != tc.wantProvID {
		t.Errorf("ProviderID = %q, want %q", got, tc.wantProvID)
	}
	if got := disp.Adapter; got != "shared-adapter" {
		t.Errorf("Adapter = %q, want %q", got, "shared-adapter")
	}
	if got := disp.Target; got != tc.wantTarget {
		t.Errorf("Target = %q, want %q", got, tc.wantTarget)
	}
	if got := disp.ResponseStallTimeoutMS; got != tc.wantTimeout {
		t.Errorf("ResponseStallTimeoutMS = %d, want %d", got, tc.wantTimeout)
	}
	if got := disp.ExecutionPath; got != tc.wantExecPath {
		t.Errorf("ExecutionPath = %q, want %q", got, tc.wantExecPath)
	}
	if got := disp.QueueReason; got != tc.wantQueueReason {
		t.Errorf("QueueReason = %q, want %q", got, tc.wantQueueReason)
	}
}

func assertTimeoutMatrixWire(t *testing.T, tc timeoutMatrixTestCase, runWires chan *iop.RunRequest, tunnelWires chan *iop.ProviderTunnelRequest) {
	runID := "run-" + tc.name
	if tc.isTunnel {
		wire := recvWire(t, tunnelWires, "ProviderTunnelRequest")
		if got := wire.GetRunId(); got != runID {
			t.Errorf("wire RunId = %q, want %q", got, runID)
		}
		if got := wire.GetTunnelId(); got != runID+"-tunnel" {
			t.Errorf("wire TunnelId = %q, want %q", got, runID+"-tunnel")
		}
		if got := wire.GetAdapter(); got != "shared-adapter" {
			t.Errorf("wire Adapter = %q, want %q", got, "shared-adapter")
		}
		if got := wire.GetTarget(); got != tc.wantTarget {
			t.Errorf("wire Target = %q, want %q", got, tc.wantTarget)
		}
		if got := wire.GetResponseStallTimeoutMs(); got != tc.wantTimeout {
			t.Errorf("wire ResponseStallTimeoutMs = %d, want %d", got, tc.wantTimeout)
		}
		assertNoExtra(t, tunnelWires, "ProviderTunnelRequest")
	} else {
		wire := recvWire(t, runWires, "RunRequest")
		if got := wire.GetRunId(); got != runID {
			t.Errorf("wire RunId = %q, want %q", got, runID)
		}
		if got := wire.GetAdapter(); got != "shared-adapter" {
			t.Errorf("wire Adapter = %q, want %q", got, "shared-adapter")
		}
		if got := wire.GetTarget(); got != tc.wantTarget {
			t.Errorf("wire Target = %q, want %q", got, tc.wantTarget)
		}
		if got := wire.GetResponseStallTimeoutMs(); got != tc.wantTimeout {
			t.Errorf("wire ResponseStallTimeoutMs = %d, want %d", got, tc.wantTimeout)
		}
		assertNoExtra(t, runWires, "RunRequest")
	}
}