package ws

import (
	"encoding/json"
	"log"
	"sync"
	"time"
)

type AgentConnection struct {
	AgentID    string
	Connection *Connection
	LastSeen   time.Time
}

type Hub struct {
	agents    map[string]*AgentConnection
	agentsMux sync.RWMutex
	
	register   chan *AgentConnection
	unregister chan *AgentConnection
	broadcast  chan []byte
	
	tunnelManager *TunnelManager

	// Response-Waiters: Command-ID -> Response-Channel
	responseWaiters map[string]chan map[string]interface{}
	waitersMux      sync.RWMutex

	// Event-Callback fuer DB-Events bei Connect/Disconnect
	eventCallback func(agentID, eventType, message string)

	// PTY-Sessions: session_id -> channel fuer Ausgabe an Browser
	ptySessions    map[string]chan []byte
	ptySessionsMux sync.RWMutex
}

func NewHub() *Hub {
	hub := &Hub{
		agents:          make(map[string]*AgentConnection),
		register:        make(chan *AgentConnection),
		unregister:      make(chan *AgentConnection),
		broadcast:       make(chan []byte, 256),
		responseWaiters: make(map[string]chan map[string]interface{}),
	}
	
	hub.tunnelManager = NewTunnelManager(hub)
	hub.ptySessions = make(map[string]chan []byte)
	return hub
}

// SetEventCallback setzt die Callback-Funktion fuer Connect/Disconnect Events
func (h *Hub) SetEventCallback(cb func(agentID, eventType, message string)) {
	h.eventCallback = cb
}

func (h *Hub) Run() {
	log.Println("WebSocket Hub startet...")
	
	// Cleanup-Timer für inaktive Verbindungen
	cleanupTicker := time.NewTicker(5 * time.Minute)
	defer cleanupTicker.Stop()
	
	for {
		select {
		case agent := <-h.register:
			h.registerAgent(agent)
			
		case agent := <-h.unregister:
			h.unregisterAgent(agent)
			
		case message := <-h.broadcast:
			h.broadcastMessage(message)
			
		case <-cleanupTicker.C:
			h.cleanupInactiveAgents()
		}
	}
}

func (h *Hub) registerAgent(agent *AgentConnection) {
	h.agentsMux.Lock()
	defer h.agentsMux.Unlock()
	
	// Bestehende Verbindung schließen falls vorhanden
	if existing, ok := h.agents[agent.AgentID]; ok {
		log.Printf("Agent %s: Ersetzt bestehende Verbindung", agent.AgentID)
		existing.Connection.Close()
	}
	
	h.agents[agent.AgentID] = agent
	agent.LastSeen = time.Now()
	
	log.Printf("Agent %s: WebSocket-Verbindung registriert", agent.AgentID)

	// Event schreiben
	if h.eventCallback != nil {
		go h.eventCallback(agent.AgentID, "connected", "WebSocket-Verbindung hergestellt")
	}
}

func (h *Hub) unregisterAgent(agent *AgentConnection) {
	h.agentsMux.Lock()
	defer h.agentsMux.Unlock()
	
	if existing, ok := h.agents[agent.AgentID]; ok && existing.Connection == agent.Connection {
		delete(h.agents, agent.AgentID)
		existing.Connection.Close()
		log.Printf("Agent %s: WebSocket-Verbindung entfernt", agent.AgentID)

		// Event schreiben
		if h.eventCallback != nil {
			go h.eventCallback(agent.AgentID, "disconnected", "WebSocket-Verbindung getrennt")
		}

		// Alle Tunnel für diesen Agent schließen
		h.tunnelManager.CloseAgentTunnels(agent.AgentID)
	}
}

func (h *Hub) broadcastMessage(message []byte) {
	h.agentsMux.RLock()
	defer h.agentsMux.RUnlock()
	
	for agentID, agent := range h.agents {
		select {
		case agent.Connection.Send <- message:
		default:
			// Agent kann keine Nachrichten mehr empfangen
			log.Printf("Agent %s: Send-Channel blockiert, trenne Verbindung", agentID)
			close(agent.Connection.Send)
			delete(h.agents, agentID)
		}
	}
}

func (h *Hub) SendToAgent(agentID string, message []byte) error {
	h.agentsMux.RLock()
	agent, exists := h.agents[agentID]
	h.agentsMux.RUnlock()
	
	if !exists {
		return ErrAgentNotConnected
	}
	
	agent.LastSeen = time.Now()
	
	select {
	case agent.Connection.Send <- message:
		return nil
	case <-time.After(5 * time.Second):
		return ErrAgentSendBlocked
	}
}

func (h *Hub) SendBinaryToAgent(agentID string, data []byte) error {
	h.agentsMux.RLock()
	agent, exists := h.agents[agentID]
	h.agentsMux.RUnlock()
	
	if !exists {
		return ErrAgentNotConnected
	}
	
	agent.LastSeen = time.Now()
	
	select {
	case agent.Connection.BinarySend <- data:
		return nil
	case <-time.After(5 * time.Second):
		return ErrAgentSendBlocked
	}
}

func (h *Hub) IsAgentConnected(agentID string) bool {
	h.agentsMux.RLock()
	defer h.agentsMux.RUnlock()
	
	_, exists := h.agents[agentID]
	return exists
}

// WaitForAgentOnline polls until the agent reconnects or timeout is reached.
// Returns true if agent came online within the timeout, false otherwise.
func (h *Hub) WaitForAgentOnline(agentID string, timeout time.Duration) bool {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if h.IsAgentConnected(agentID) {
			return true
		}
		time.Sleep(10 * time.Second)
	}
	return false
}

func (h *Hub) GetConnectedAgents() []string {
	h.agentsMux.RLock()
	defer h.agentsMux.RUnlock()
	
	agents := make([]string, 0, len(h.agents))
	for agentID := range h.agents {
		agents = append(agents, agentID)
	}
	
	return agents
}

func (h *Hub) GetAgentStats() map[string]interface{} {
	h.agentsMux.RLock()
	defer h.agentsMux.RUnlock()
	
	stats := map[string]interface{}{
		"connected_agents": len(h.agents),
		"active_tunnels":   h.tunnelManager.GetActiveTunnelCount(),
	}
	
	return stats
}

func (h *Hub) cleanupInactiveAgents() {
	h.agentsMux.Lock()
	defer h.agentsMux.Unlock()
	
	timeout := 10 * time.Minute
	now := time.Now()
	
	for agentID, agent := range h.agents {
		if now.Sub(agent.LastSeen) > timeout {
			log.Printf("Agent %s: Timeout, trenne Verbindung", agentID)
			delete(h.agents, agentID)
			agent.Connection.Close()
			h.tunnelManager.CloseAgentTunnels(agentID)
		}
	}
}

// RegisterResponseWaiter registriert einen Channel fuer eine Command-Response
func (h *Hub) RegisterResponseWaiter(cmdID string) chan map[string]interface{} {
	ch := make(chan map[string]interface{}, 1)
	h.waitersMux.Lock()
	h.responseWaiters[cmdID] = ch
	h.waitersMux.Unlock()
	return ch
}

// UnregisterResponseWaiter entfernt einen Response-Waiter
func (h *Hub) UnregisterResponseWaiter(cmdID string) {
	h.waitersMux.Lock()
	delete(h.responseWaiters, cmdID)
	h.waitersMux.Unlock()
}

// GetTunnelManager gibt den TunnelManager zurück (für API-Handler)
func (h *Hub) GetTunnelManager() *TunnelManager {
	return h.tunnelManager
}

// RegisterAgent registriert eine neue Agent-Verbindung
func (h *Hub) RegisterAgent(agentID string, conn *Connection) {
	agent := &AgentConnection{
		AgentID:    agentID,
		Connection: conn,
		LastSeen:   time.Now(),
	}
	
	h.register <- agent
}

// UnregisterAgent entfernt eine Agent-Verbindung
func (h *Hub) UnregisterAgent(agentID string, conn *Connection) {
	agent := &AgentConnection{
		AgentID:    agentID,
		Connection: conn,
	}
	
	h.unregister <- agent
}

// ProcessAgentMessage verarbeitet eingehende Nachrichten von Agents
func (h *Hub) ProcessAgentMessage(agentID string, messageType int, data []byte) error {
	h.agentsMux.RLock()
	agent, exists := h.agents[agentID]
	h.agentsMux.RUnlock()
	
	if exists {
		agent.LastSeen = time.Now()
	}
	
	// Binary Messages: PTY oder Tunnel
	if messageType == 2 { // websocket.BinaryMessage
		// PTY-Session erkennen: session_id beginnt mit "pty-"
		if len(data) > 1 {
			idLen := int(data[0])
			if 1+idLen <= len(data) {
				sessionID := string(data[1 : 1+idLen])
				payload := data[1+idLen:]
				if len(sessionID) >= 4 && sessionID[:4] == "pty-" {
					h.ptySessionsMux.RLock()
					ch, ok := h.ptySessions[sessionID]
					h.ptySessionsMux.RUnlock()
					if ok {
						select {
						case ch <- payload:
						default:
						}
					}
					return nil
				}
			}
		}
		return h.tunnelManager.ProcessBinaryMessage(agentID, data)
	}
	
	// Text Messages sind JSON-Responses/Events vom Agent
	var msg struct {
		Type   string                 `json:"type"`
		ID     string                 `json:"id"`
		Status string                 `json:"status"`
		Data   map[string]interface{} `json:"data"`
	}
	if err := json.Unmarshal(data, &msg); err != nil {
		return nil // Ignorieren wenn nicht parsebar
	}

	if msg.Type == "response" {
		// Tunnel-Connect Responses verarbeiten
		if msg.Status == "ok" && msg.Data != nil {
			if sessionID, ok := msg.Data["session_id"].(string); ok {
				h.tunnelManager.ConfirmSession(sessionID)
			}
		}

		// Response an wartende API-Handler weiterleiten
		if msg.ID != "" {
			h.waitersMux.RLock()
			ch, exists := h.responseWaiters[msg.ID]
			h.waitersMux.RUnlock()

			if exists {
				resp := map[string]interface{}{
					"status": msg.Status,
					"data":   msg.Data,
				}
				if msg.Status == "error" {
					var errMsg struct {
						Error string `json:"error"`
					}
					json.Unmarshal(data, &errMsg)
					resp["error"] = errMsg.Error
				}

				select {
				case ch <- resp:
				default:
				}
			}
		}
	}

	return nil
}
// RegisterPTYSession registriert einen Output-Channel fuer eine PTY-Session
func (h *Hub) RegisterPTYSession(sessionID string) chan []byte {
	ch := make(chan []byte, 256)
	h.ptySessionsMux.Lock()
	h.ptySessions[sessionID] = ch
	h.ptySessionsMux.Unlock()
	return ch
}

// UnregisterPTYSession entfernt eine PTY-Session
func (h *Hub) UnregisterPTYSession(sessionID string) {
	h.ptySessionsMux.Lock()
	delete(h.ptySessions, sessionID)
	h.ptySessionsMux.Unlock()
}