from fastapi import WebSocket
from typing import Dict
import asyncio

class ConnectionManager:
    """Manages all active WebSocket connections"""

    def __init__(self):
        self.active: Dict[str, WebSocket] = {}  # client_id → socket

    async def connect(self, ws: WebSocket, client_id: str):
        await ws.accept()
        self.active[client_id] = ws
        await self.broadcast({
            "type": "user_joined",
            "client_id": client_id,
            "online_count": len(self.active)
        }, exclude=client_id)

    def disconnect(self, client_id: str):
        self.active.pop(client_id, None)

    async def send_to(self, client_id: str, message: dict):
        """Send to specific client"""
        ws = self.active.get(client_id)
        if ws:
            await ws.send_json(message)

    async def broadcast(self, message: dict, exclude: str = None):
        """Send to ALL connected clients"""
        dead = []
        for cid, ws in self.active.items():
            if cid == exclude: continue
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(cid)  # connection died silently
        for cid in dead:
            self.disconnect(cid)

# Singleton instance
manager = ConnectionManager()

@router.websocket("/ws/chat/{client_id}")
async def chat_ws(websocket: WebSocket, client_id: str):
    await manager.connect(websocket, client_id)
    try:
        while True:
            data = await websocket.receive_json()
            await manager.broadcast({
                "type": "message",
                "from": client_id,
                "text": data.get("text"),
                "timestamp": datetime.utcnow().isoformat()
            })
    except WebSocketDisconnect:
        manager.disconnect(client_id)
        await manager.broadcast({
            "type": "user_left",
            "client_id": client_id
        })

from collections import defaultdict

class RoomManager:
    def __init__(self):
        # room_id → {client_id: websocket}
        self.rooms: Dict[str, Dict[str, WebSocket]] = defaultdict(dict)

    async def join(self, ws: WebSocket, room_id: str, client_id: str):
        await ws.accept()
        self.rooms[room_id][client_id] = ws
        await self.broadcast_to_room(room_id, {
            "type": "joined", "user": client_id,
            "members": list(self.rooms[room_id].keys())
        })

    async def leave(self, room_id: str, client_id: str):
        self.rooms[room_id].pop(client_id, None)
        if not self.rooms[room_id]:
            del self.rooms[room_id]  # empty room cleanup
        else:
            await self.broadcast_to_room(room_id, {
                "type": "left", "user": client_id
            })

    async def broadcast_to_room(self, room_id: str, message: dict):
        for ws in self.rooms.get(room_id, {}).values():
            try:
                await ws.send_json(message)
            except: pass

room_manager = RoomManager()

@router.websocket("/ws/room/{room_id}/{client_id}")
async def room_ws(ws: WebSocket, room_id: str, client_id: str):
    await room_manager.join(ws, room_id, client_id)
    try:
        while True:
            data = await ws.receive_json()
            await room_manager.broadcast_to_room(room_id, {
                "type": "message", "from": client_id,
                "text": data.get("text")
            })
    except WebSocketDisconnect:
        await room_manager.leave(room_id, client_id)

import { Component, OnInit, OnDestroy, inject } from '@angular/core';
import { Subject } from 'rxjs';
import { takeUntil } from 'rxjs/operators';
import { WebSocketService } from '../services/websocket.service';

@Component({
  selector: 'app-chat',
  template: `
    <div *ngFor="let msg of messages">
      <strong>{{ msg.from }}:</strong> {{ msg.text }}
    </div>
    <input [(ngModel)]="newMessage" (keyup.enter)="send()" />
    <button (click)="send()">Send</button>
    <span>{{ onlineCount }} online</span>
  `
})
export class ChatComponent implements OnInit, OnDestroy {
  private ws = inject(WebSocketService);
  private destroy$ = new Subject<void>();

  messages: any[] = [];
  newMessage = '';
  onlineCount = 0;
  clientId = crypto.randomUUID().slice(0, 8);

  ngOnInit(): void {
    // Connect to WebSocket
    this.ws.connect(`ws://localhost:8000/ws/chat/${this.clientId}`);

    // Listen for chat messages
    this.ws.on('message').pipe(
      takeUntil(this.destroy$)
    ).subscribe(msg => {
      this.messages.push(msg);
    });

    // Listen for user joins
    this.ws.on('user_joined').pipe(
      takeUntil(this.destroy$)
    ).subscribe(msg => {
      this.onlineCount = msg.online_count;
    });

    // Or subscribe to ALL messages
    this.ws.messages$.pipe(
      takeUntil(this.destroy$)
    ).subscribe(msg => {
      switch (msg.type) {
        case 'message': this.messages.push(msg); break;
        case 'user_joined': this.onlineCount = msg.online_count; break;
        case 'user_left': this.onlineCount--; break;
      }
    });
  }

  send(): void {
    if (!this.newMessage.trim()) return;
    this.ws.send({ type: 'message', text: this.newMessage });
    this.newMessage = '';
  }

  ngOnDestroy(): void {
    this.destroy$.next();    // unsubscribe from all observables
    this.destroy$.complete();
    this.ws.disconnect();
  }
}

@Component({ selector: 'app-chat', templateUrl: './chat.component.html' })
export class ChatComponent implements OnInit, OnDestroy {
  private chat = inject(ChatService);
  private destroy$ = new Subject<void>();

  messages: ChatMessage[] = [];
  newMsg = '';
  username = '';
  joined = false;
  typingUser = '';
  onlineUsers: string[] = [];
  isConnected$ = this.chat.isConnected$;
  onlineUsers$ = this.chat.onlineUsers$;

  ngOnInit(): void {
    this.chat.messages$.pipe(
      takeUntil(this.destroy$)
    ).subscribe(msg => {
      switch (msg.type) {
        case 'message':
        case 'system':
          this.messages.push(msg);
          break;
        case 'history':
          this.messages = (msg as any).messages;
          break;
        case 'typing':
          this.typingUser = msg.from!;
          setTimeout(() => this.typingUser = '', 2000);
          break;
      }
    });
  }

  join(): void {
    if (!this.username.trim()) return;
    this.chat.connect(this.username);
    this.joined = true;
  }

  send(): void {
    if (!this.newMsg.trim()) return;
    this.chat.sendMessage(this.newMsg);
    this.newMsg = '';
  }

  onTyping(): void { this.chat.sendTyping(); }

  ngOnDestroy(): void { this.destroy$.next(); }
}