package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"sync"
	"testing"

	"code.cubecraftcreations.com/CubeCraft-Creations/Control-Center/go-backend/internal/handler"
	"code.cubecraftcreations.com/CubeCraft-Creations/Control-Center/go-backend/internal/models"
)

// ── Mock AgentRepo ────────────────────────────────────────────────────────

type mockAgentRepo struct {
	mu          sync.Mutex
	agents      map[string]models.AgentCardData
	updateCalls []updateCall
}

type updateCall struct {
	id  string
	req models.UpdateAgentRequest
}

func (m *mockAgentRepo) Get(_ context.Context, id string) (models.AgentCardData, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	a, ok := m.agents[id]
	if !ok {
		return models.AgentCardData{}, errNotFound
	}
	return a, nil
}

func (m *mockAgentRepo) Update(_ context.Context, id string, req models.UpdateAgentRequest) (models.AgentCardData, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	a, ok := m.agents[id]
	if !ok {
		return models.AgentCardData{}, errNotFound
	}

	if req.Status != nil {
		a.Status = *req.Status
	}
	if req.DisplayName != nil {
		a.DisplayName = *req.DisplayName
	}
	if req.Role != nil {
		a.Role = *req.Role
	}
	if req.Channel != nil {
		a.Channel = *req.Channel
	}
	if req.CurrentTask != nil {
		a.CurrentTask = req.CurrentTask
	}
	if req.TaskProgress != nil {
		a.TaskProgress = req.TaskProgress
	}
	if req.TaskElapsed != nil {
		a.TaskElapsed = req.TaskElapsed
	}
	if req.ErrorMessage != nil {
		a.ErrorMessage = req.ErrorMessage
	}
	if req.LastActivityAt != nil {
		a.LastActivity = *req.LastActivityAt
	}

	m.agents[id] = a
	m.updateCalls = append(m.updateCalls, updateCall{id, req})
	return a, nil
}

func (m *mockAgentRepo) Create(_ context.Context, a models.AgentCardData) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.agents[a.ID] = a
	return nil
}

func (m *mockAgentRepo) List(_ context.Context, statusFilter models.AgentStatus) ([]models.AgentCardData, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	var result []models.AgentCardData
	for _, a := range m.agents {
		if statusFilter == "" || a.Status == statusFilter {
			result = append(result, a)
		}
	}
	return result, nil
}

func (m *mockAgentRepo) Delete(_ context.Context, id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.agents, id)
	return nil
}

func (m *mockAgentRepo) Count(_ context.Context) (int, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return len(m.agents), nil
}

// errNotFound is returned by the mock repo when an agent is not found.
var errNotFound = fmt.Errorf("not found")

// ── Broadcast capture helper ───────────────────────────────────────────────

// broadcastCapture wraps a real Broker and captures all broadcasts
// via a subscribed channel. Use captured() to retrieve events that have
// been received so far. Call close() to unsubscribe when done.
type broadcastCapture struct {
	broker *handler.Broker
	ch     chan handler.SSEEvent
}

func newBroadcastCapture(broker *handler.Broker) *broadcastCapture {
	return &broadcastCapture{
		broker: broker,
		ch:     broker.Subscribe(),
	}
}

// captured drains all pending events from the subscription channel
// and returns them. This is synchronous — it only returns events that
// have already been sent to the channel.
func (bc *broadcastCapture) captured() []handler.SSEEvent {
	var events []handler.SSEEvent
	for {
		select {
		case evt := <-bc.ch:
			events = append(events, evt)
		default:
			return events
		}
	}
}

func (bc *broadcastCapture) close() {
	bc.broker.Unsubscribe(bc.ch)
}

// ── Test helpers ──────────────────────────────────────────────────────────

// newTestWSClient creates a WSClient wired to a mock repo and a real broker.
// Returns the client, the mock repo, and a broadcast capture.
func newTestWSClient() (*WSClient, *mockAgentRepo, *handler.Broker, *broadcastCapture) {
	repo := &mockAgentRepo{agents: make(map[string]models.AgentCardData)}
	broker := handler.NewBroker()
	capture := newBroadcastCapture(broker)
	client := NewWSClient(WSConfig{}, repo, broker, slog.Default())
	return client, repo, broker, capture
}

// ── Tests ─────────────────────────────────────────────────────────────────

func TestHandleSessionsChanged_Active(t *testing.T) {
	client, repo, broker, capture := newTestWSClient()
	defer capture.close()

	repo.agents["otto"] = models.AgentCardData{
		ID:          "otto",
		DisplayName: "Otto",
		Status:      models.AgentStatusIdle,
	}

	payload := json.RawMessage(`{
		"sessionKey": "s1",
		"agentId":    "otto",
		"status":     "running",
		"totalTokens": 500,
		"currentTask": "Orchestrating tasks"
	}`)

	client.handleSessionsChanged(payload)

	// Verify: agent status updated to active
	repo.mu.Lock()
	agent := repo.agents["otto"]
	calls := make([]updateCall, len(repo.updateCalls))
	copy(calls, repo.updateCalls)
	repo.mu.Unlock()

	if agent.Status != models.AgentStatusActive {
		t.Errorf("agent status = %q, want %q", agent.Status, models.AgentStatusActive)
	}

	// Verify: update was called
	if len(calls) == 0 {
		t.Fatal("expected at least one update call")
	}
	if calls[0].id != "otto" {
		t.Errorf("update call agentId = %q, want %q", calls[0].id, "otto")
	}

	// Verify: broker broadcast "agent.status"
	events := capture.captured()
	found := false
	for _, evt := range events {
		if evt.EventType == "agent.status" {
			found = true
			break
		}
	}
	if !found {
		t.Error("expected broker broadcast with event type 'agent.status'")
	}
}

func TestHandleSessionsChanged_Idle(t *testing.T) {
	client, repo, _, capture := newTestWSClient()
	defer capture.close()

	repo.agents["dex"] = models.AgentCardData{
		ID:          "dex",
		DisplayName: "Dex",
		Status:      models.AgentStatusActive,
		CurrentTask: strPtr("Writing API"),
	}

	payload := json.RawMessage(`{
		"sessionKey": "s2",
		"agentId":    "dex",
		"status":     "done",
		"totalTokens": 1000
	}`)

	client.handleSessionsChanged(payload)

	repo.mu.Lock()
	agent := repo.agents["dex"]
	repo.mu.Unlock()

	// Verify: agent goes idle
	if agent.Status != models.AgentStatusIdle {
		t.Errorf("agent status = %q, want %q", agent.Status, models.AgentStatusIdle)
	}

	// Verify: current task cleared (set to empty string)
	if agent.CurrentTask != nil && *agent.CurrentTask != "" {
		t.Errorf("current task = %q, want empty (cleared on idle)", *agent.CurrentTask)
	}

	// Verify: broker fires "agent.status"
	events := capture.captured()
	found := false
	for _, evt := range events {
		if evt.EventType == "agent.status" {
			found = true
			break
		}
	}
	if !found {
		t.Error("expected broker broadcast with event type 'agent.status'")
	}
}

func TestHandleSessionsChanged_ArrayPayload(t *testing.T) {
	client, repo, _, capture := newTestWSClient()
	defer capture.close()

	repo.agents["otto"] = models.AgentCardData{ID: "otto", DisplayName: "Otto", Status: models.AgentStatusIdle}
	repo.agents["dex"] = models.AgentCardData{ID: "dex", DisplayName: "Dex", Status: models.AgentStatusIdle}

	payload := json.RawMessage(`[
		{"sessionKey":"s1","agentId":"otto","status":"running","totalTokens":100},
		{"sessionKey":"s2","agentId":"dex","status":"streaming","totalTokens":200}
	]`)

	client.handleSessionsChanged(payload)

	repo.mu.Lock()
	otto := repo.agents["otto"]
	dex := repo.agents["dex"]
	repo.mu.Unlock()

	if otto.Status != models.AgentStatusActive {
		t.Errorf("otto status = %q, want active", otto.Status)
	}
	if dex.Status != models.AgentStatusActive {
		t.Errorf("dex status = %q, want active", dex.Status)
	}

	// Both should produce broadcasts
	events := capture.captured()
	statusCount := 0
	for _, evt := range events {
		if evt.EventType == "agent.status" {
			statusCount++
		}
	}
	if statusCount < 2 {
		t.Errorf("expected at least 2 agent.status broadcasts, got %d", statusCount)
	}
}

func TestHandleSessionsChanged_SkipsEmptyAgentID(t *testing.T) {
	client, _, _, capture := newTestWSClient()
	defer capture.close()

	payload := json.RawMessage(`{"sessionKey":"s1","agentId":"","status":"running"}`)
	client.handleSessionsChanged(payload)

	events := capture.captured()
	if len(events) > 0 {
		t.Errorf("expected no broadcasts for empty agentId, got %d", len(events))
	}
}

func TestHandleSessionsChanged_UnparseablePayload(t *testing.T) {
	client, _, _, capture := newTestWSClient()
	defer capture.close()

	payload := json.RawMessage(`not json at all`)
	client.handleSessionsChanged(payload)

	events := capture.captured()
	if len(events) > 0 {
		t.Errorf("expected no broadcasts for unparseable payload, got %d", len(events))
	}
}

func TestHandlePresence(t *testing.T) {
	client, repo, _, capture := newTestWSClient()
	defer capture.close()

	repo.agents["pip"] = models.AgentCardData{
		ID:          "pip",
		DisplayName: "Pip",
		Status:      models.AgentStatusActive,
	}

	payload := json.RawMessage(`{
		"agentId":        "pip",
		"connected":      true,
		"lastActivityAt": "2025-01-01T00:00:00Z"
	}`)

	client.handlePresence(payload)

	repo.mu.Lock()
	agent := repo.agents["pip"]
	calls := make([]updateCall, len(repo.updateCalls))
	copy(calls, repo.updateCalls)
	repo.mu.Unlock()

	// Agent should still be active (connected=true doesn't change status)
	if agent.Status != models.AgentStatusActive {
		t.Errorf("agent status = %q, want active", agent.Status)
	}

	// Update should have been called (for lastActivityAt)
	if len(calls) == 0 {
		t.Fatal("expected at least one update call")
	}

	// Verify broadcast
	events := capture.captured()
	found := false
	for _, evt := range events {
		if evt.EventType == "agent.status" {
			found = true
			break
		}
	}
	if !found {
		t.Error("expected broker broadcast with event type 'agent.status'")
	}
}

func TestHandlePresence_Disconnect(t *testing.T) {
	client, repo, _, capture := newTestWSClient()
	defer capture.close()

	repo.agents["pip"] = models.AgentCardData{
		ID:          "pip",
		DisplayName: "Pip",
		Status:      models.AgentStatusActive,
	}

	payload := json.RawMessage(`{
		"agentId":   "pip",
		"connected": false
	}`)

	client.handlePresence(payload)

	repo.mu.Lock()
	agent := repo.agents["pip"]
	repo.mu.Unlock()

	// Agent should go idle on disconnect
	if agent.Status != models.AgentStatusIdle {
		t.Errorf("agent status = %q, want idle after disconnect", agent.Status)
	}

	events := capture.captured()
	found := false
	for _, evt := range events {
		if evt.EventType == "agent.status" {
			found = true
			break
		}
	}
	if !found {
		t.Error("expected broker broadcast with event type 'agent.status' on disconnect")
	}
}

func TestHandlePresence_EmptyAgentID(t *testing.T) {
	client, _, _, capture := newTestWSClient()
	defer capture.close()

	payload := json.RawMessage(`{"agentId":"","connected":true}`)
	client.handlePresence(payload)

	events := capture.captured()
	if len(events) > 0 {
		t.Errorf("expected no broadcasts for empty agentId, got %d", len(events))
	}
}

func TestHandleAgentConfig(t *testing.T) {
	client, repo, _, capture := newTestWSClient()
	defer capture.close()

	repo.agents["rex"] = models.AgentCardData{
		ID:          "rex",
		DisplayName:  "Rex",
		Role:        "Frontend Dev",
		Status:      models.AgentStatusIdle,
		Channel:     "discord",
	}

	payload := json.RawMessage(`{
		"id":      "rex",
		"name":    "Rex the Dev",
		"role":    "Senior Frontend",
		"channel": "telegram"
	}`)

	client.handleAgentConfig(payload)

	repo.mu.Lock()
	agent := repo.agents["rex"]
	calls := make([]updateCall, len(repo.updateCalls))
	copy(calls, repo.updateCalls)
	repo.mu.Unlock()

	// Verify DisplayName and Role updated
	if agent.DisplayName != "Rex the Dev" {
		t.Errorf("displayName = %q, want %q", agent.DisplayName, "Rex the Dev")
	}
	if agent.Role != "Senior Frontend" {
		t.Errorf("role = %q, want %q", agent.Role, "Senior Frontend")
	}
	if agent.Channel != "telegram" {
		t.Errorf("channel = %q, want %q", agent.Channel, "telegram")
	}

	// Verify update was called
	if len(calls) == 0 {
		t.Fatal("expected at least one update call")
	}

	// Verify broker fires "fleet.update"
	events := capture.captured()
	found := false
	for _, evt := range events {
		if evt.EventType == "fleet.update" {
			found = true
			break
		}
	}
	if !found {
		t.Error("expected broker broadcast with event type 'fleet.update'")
	}
}

func TestHandleAgentConfig_EmptyID(t *testing.T) {
	client, _, _, capture := newTestWSClient()
	defer capture.close()

	payload := json.RawMessage(`{"id":"","name":"Ghost"}`)
	client.handleAgentConfig(payload)

	events := capture.captured()
	if len(events) > 0 {
		t.Errorf("expected no broadcasts for empty id, got %d", len(events))
	}
}

func TestHandleAgentConfig_NotFound(t *testing.T) {
	client, _, _, capture := newTestWSClient()
	defer capture.close()

	payload := json.RawMessage(`{"id":"unknown","name":"Ghost","role":"Phantom"}`)
	client.handleAgentConfig(payload)

	// Agent doesn't exist in repo, so Update will fail → handler logs warning, returns early
	events := capture.captured()
	for _, evt := range events {
		if evt.EventType == "fleet.update" {
			t.Error("fleet.update should not be broadcast when agent update fails")
		}
	}
}