package main

import (
	"bytes"
	"context"
	"strings"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/common-proto-socket/go"
	edgenode "iop/apps/edge/internal/node"
	iop "iop/proto/gen/iop"
)

func TestBuildRunRequest_SessionAndBackground(t *testing.T) {
	req, runID, err := buildRunRequest("cli", "codex", "session-a", true, 30, "hello")
	if err != nil {
		t.Fatalf("buildRunRequest: %v", err)
	}
	if runID == "" {
		t.Fatal("expected non-empty runID")
	}
	if req.GetAdapter() != "cli" {
		t.Errorf("Adapter: got %q want %q", req.GetAdapter(), "cli")
	}
	if req.GetModel() != "codex" {
		t.Errorf("wire model: got %q want %q", req.GetModel(), "codex")
	}
	if req.GetSessionId() != "session-a" {
		t.Errorf("SessionId: got %q want %q", req.GetSessionId(), "session-a")
	}
	if !req.GetBackground() {
		t.Error("Background: expected true")
	}
	if req.GetSessionMode() != iop.RunSessionMode_RUN_SESSION_MODE_CREATE_IF_MISSING {
		t.Errorf("SessionMode: got %v", req.GetSessionMode())
	}
}

func TestBuildRunRequest_EmptySessionNormalized(t *testing.T) {
	req, _, err := buildRunRequest("cli", "codex", "", false, 30, "hello")
	if err != nil {
		t.Fatalf("buildRunRequest: %v", err)
	}
	if req.GetSessionId() != "default" {
		t.Errorf("SessionId: got %q want %q", req.GetSessionId(), "default")
	}
}

func TestBuildRunRequest_DefaultTimeoutFallback(t *testing.T) {
	req, _, err := buildRunRequest("cli", "codex", "s1", false, 0, "hello")
	if err != nil {
		t.Fatalf("buildRunRequest: %v", err)
	}
	if req.GetTimeoutSec() != 30 {
		t.Errorf("TimeoutSec: got %d want 30", req.GetTimeoutSec())
	}
}

func TestNormalizeConsoleSessionID(t *testing.T) {
	if got := normalizeConsoleSessionID(""); got != "default" {
		t.Errorf("empty: got %q want %q", got, "default")
	}
	if got := normalizeConsoleSessionID("my-session"); got != "my-session" {
		t.Errorf("non-empty: got %q want %q", got, "my-session")
	}
}

func TestConsoleEventRouterRoutesRegisteredRun(t *testing.T) {
	var out bytes.Buffer
	router := newConsoleEventRouter(&out, nil, nil)
	events, unregister := router.Register("run-1")
	defer unregister()

	router.Handle(&iop.RunEvent{RunId: "run-1", Type: "start"})

	select {
	case event := <-events:
		if event.GetRunId() != "run-1" {
			t.Fatalf("RunId: got %q want run-1", event.GetRunId())
		}
	default:
		t.Fatal("expected registered run event")
	}
	if out.Len() != 0 {
		t.Fatalf("registered foreground event should not be printed asynchronously: %q", out.String())
	}
}

func TestConsoleEventRouterPrintsUnregisteredRun(t *testing.T) {
	var out bytes.Buffer
	router := newConsoleEventRouter(&out, nil, nil)

	router.Handle(&iop.RunEvent{RunId: "run-bg", Type: "start", SessionId: "s1", Background: true})
	router.Handle(&iop.RunEvent{RunId: "run-bg", Type: "delta", Delta: "hello"})
	if got := out.String(); !strings.Contains(got, "[node-unknown-message] hello") {
		t.Fatalf("expected delta to be printed before completion, got:\n%s", got)
	}
	router.Handle(&iop.RunEvent{RunId: "run-bg", Type: "delta", Delta: " world\n"})
	router.Handle(&iop.RunEvent{RunId: "run-bg", Type: "complete", Message: "done"})

	got := out.String()
	for _, want := range []string{
		"[node-unknown-event] start run_id=run-bg session=s1 background=true",
		"[node-unknown-event] complete run_id=run-bg detail=\"done\"",
		"[node-unknown-message] hello world",
	} {
		if !strings.Contains(got, want) {
			t.Fatalf("expected output to contain %q, got:\n%s", want, got)
		}
	}
}

func TestConsoleResponseStreamWritesBeforeFinish(t *testing.T) {
	var out bytes.Buffer
	stream := newConsoleResponseStream(&out, "[node-test-message] ")

	stream.Write("hello")
	if got := out.String(); got != "[node-test-message] hello" {
		t.Fatalf("Write should print immediately, got %q", got)
	}

	stream.Write(" world")
	stream.Finish()
	if got := out.String(); got != "[node-test-message] hello world\n" {
		t.Fatalf("Finish should close the line, got %q", got)
	}
}

func TestBuildNodeCommandRequest_StatusUsesCurrentTarget(t *testing.T) {
	req, reqID := buildNodeCommandRequest("cli", "codex", "default", 30)
	if reqID == "" {
		t.Fatal("expected non-empty reqID")
	}
	if req.GetType() != iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS {
		t.Errorf("Type: got %v want %v", req.GetType(), iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS)
	}
	if req.GetAdapter() != "cli" {
		t.Errorf("Adapter: got %q want %q", req.GetAdapter(), "cli")
	}
	if req.GetModel() != "codex" {
		t.Errorf("Model: got %q want %q", req.GetModel(), "codex")
	}
	if req.GetSessionId() != "default" {
		t.Errorf("SessionId: got %q want %q", req.GetSessionId(), "default")
	}
	if req.GetTimeoutSec() != 30 {
		t.Errorf("TimeoutSec: got %d want 30", req.GetTimeoutSec())
	}
}

func TestStatusWaitTimeout_UsesNormalizedRequestTimeout(t *testing.T) {
	// 0 should be normalized to 30, and then +5 added
	req, _ := buildNodeCommandRequest("cli", "codex", "default", 0)
	got := statusWaitTimeout(req)
	want := 35 * time.Second
	if got != want {
		t.Errorf("expected %v, got %v", want, got)
	}

	// Explicit 60 should result in 65
	req2, _ := buildNodeCommandRequest("cli", "codex", "default", 60)
	got2 := statusWaitTimeout(req2)
	want2 := 65 * time.Second
	if got2 != want2 {
		t.Errorf("expected %v, got %v", want2, got2)
	}
}

func TestFormatUsageStatus_PrintsParsedLimits(t *testing.T) {
	var out bytes.Buffer
	formatUsageStatus(&out, "local-node", "codex", "default", &iop.AgentUsageStatus{
		DailyLimit:  "98% left",
		WeeklyLimit: "22% left",
	})
	got := out.String()
	if !strings.Contains(got, "Daily limit: 98% left") {
		t.Errorf("expected Daily limit: 98%% left, got:\n%s", got)
	}
	if !strings.Contains(got, "Weekly limit: 22% left") {
		t.Errorf("expected Weekly limit: 22%% left, got:\n%s", got)
	}
}

func TestFormatUsageStatus_Error(t *testing.T) {
	// The node reports an error string directly inside the command response, which is checked before formatUsageStatus is called.
	// But we should test if raw status block is printed properly if no limits are parsed.
	var out bytes.Buffer
	formatUsageStatus(&out, "local-node", "codex", "default", &iop.AgentUsageStatus{
		RawOutput: "unparsed error text",
	})
	got := out.String()
	if !strings.Contains(got, "raw output did not include parsed limits:") {
		t.Errorf("expected raw output warning, got:\n%s", got)
	}
	if !strings.Contains(got, "unparsed error text") {
		t.Errorf("expected unparsed text, got:\n%s", got)
	}
}

func TestResolveConsoleNode_RequiresExplicitSelectionForMultipleNodes(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1"})
	reg.Register(&edgenode.NodeEntry{NodeID: "node-2"})

	// Empty target should fail for multiple nodes
	if _, err := resolveConsoleNode(reg, ""); err == nil {
		t.Error("expected error for implicit resolve with multiple nodes")
	}

	// Explicit target should succeed
	if e, err := resolveConsoleNode(reg, "node-1"); err != nil || e.NodeID != "node-1" {
		t.Errorf("failed explicit resolve: %v", err)
	}
}

func TestResolveConsoleNode_AllowsSingleNodeFallback(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1"})

	if e, err := resolveConsoleNode(reg, ""); err != nil || e.NodeID != "node-1" {
		t.Errorf("expected single node fallback, got error: %v", err)
	}
}

func TestPrintNodes_ShowsSelectedNode(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1"})
	reg.Register(&edgenode.NodeEntry{NodeID: "node-2", Alias: "alias-2"})

	var out bytes.Buffer
	printNodes(&out, reg, "alias-1")

	got := out.String()
	if !strings.Contains(got, "* node-1 (alias-1)") {
		t.Errorf("expected selected node-1 to have marker, got:\n%s", got)
	}
	if !strings.Contains(got, "  node-2 (alias-2)") {
		t.Errorf("expected non-selected node-2 to have no marker, got:\n%s", got)
	}
}

func TestSendTerminateSession_ReturnsNodeLabel(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID: "node-1",
		Alias:  "alias-1",
		Client: &toki.TcpClient{}, // mock client
	})

	target := &consoleTarget{
		NodeRef:   "alias-1",
		SessionID: "s1",
	}

	label, _ := sendTerminateSession(context.Background(), reg, target)
	if label != "alias-1" {
		t.Errorf("expected label alias-1, got %q", label)
	}

	// Test fallback to NodeID
	reg2 := edgenode.NewRegistry()
	reg2.Register(&edgenode.NodeEntry{
		NodeID: "node-2",
		Client: &toki.TcpClient{},
	})
	target.NodeRef = "node-2"
	label2, _ := sendTerminateSession(context.Background(), reg2, target)
	if label2 != "node-2" {
		t.Errorf("expected label node-2, got %q", label2)
	}
}

func TestHandleTerminateSession_OutputFormat(t *testing.T) {
	// Mock sendTerminateSessionFunc to avoid network/toki dependency
	oldFunc := sendTerminateSessionFunc
	defer func() { sendTerminateSessionFunc = oldFunc }()

	sendTerminateSessionFunc = func(ctx context.Context, registry *edgenode.Registry, target *consoleTarget) (string, error) {
		return "mock-alias", nil
	}

	reg := edgenode.NewRegistry()
	target := &consoleTarget{
		SessionID: "s-test",
	}

	var out bytes.Buffer
	handleTerminateSession(context.Background(), reg, &out, target)

	got := out.String()
	want := "terminated session s-test node=mock-alias\n"
	if got != want {
		t.Errorf("OutputFormat (alias): got %q, want %q", got, want)
	}

	// Test fallback to NodeID in output
	sendTerminateSessionFunc = func(ctx context.Context, registry *edgenode.Registry, target *consoleTarget) (string, error) {
		return "node-raw-id", nil
	}
	out.Reset()
	handleTerminateSession(context.Background(), reg, &out, target)
	got2 := out.String()
	want2 := "terminated session s-test node=node-raw-id\n"
	if got2 != want2 {
		t.Errorf("OutputFormat (node_id): got %q, want %q", got2, want2)
	}
}

func TestConsoleEventRouterPrintsNodeScopedAsyncRun(t *testing.T) {
	var out bytes.Buffer
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1"})
	router := newConsoleEventRouter(&out, reg, nil)

	router.Handle(&iop.RunEvent{
		RunId:      "run-1",
		Type:       "start",
		NodeId:     "node-1",
		Background: true,
	})

	got := out.String()
	if !strings.Contains(got, "[node-alias-1-event] start run_id=run-1") {
		t.Errorf("expected node-alias-1-event label, got:\n%s", got)
	}
}

func TestConsoleEventRouterFallsBackToNodeID(t *testing.T) {
	var out bytes.Buffer
	router := newConsoleEventRouter(&out, nil, nil)

	router.Handle(&iop.RunEvent{
		RunId:      "run-1",
		Type:       "start",
		NodeId:     "some-raw-id",
		Background: true,
	})

	got := out.String()
	if !strings.Contains(got, "[node-some-raw-id-event] start run_id=run-1") {
		t.Errorf("expected node-some-raw-id-event label, got:\n%s", got)
	}
}