zyrarbot/web.py

from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, UploadFile, File
from pydantic import BaseModel
from typing import List, Optional
import asyncio
import discord
import os
import shutil

# Import bot instance and YTDL class
from bot import bot
from ytdl import YTDLSource

# Initialize FastAPI application.
# The service exposes the API only - the web panel is a separate container.
app = FastAPI(title="ŻyrarBOT API")

# FastAPI Request Models
class PlayRequest(BaseModel):
    query: str

class VolumeRequest(BaseModel):
    volume: int

class LoopRequest(BaseModel):
    loop: bool

class RemoveRequest(BaseModel):
    id: str

class ReorderRequest(BaseModel):
    queue_ids: List[str]

class JumpRequest(BaseModel):
    index: int

class ConnectRequest(BaseModel):
    channel_id: str

class TextChannelRequest(BaseModel):
    channel_id: str
    guild_id: Optional[str] = None

class PlayFileRequest(BaseModel):
    filename: str

class DefaultPlaysRequest(BaseModel):
    default_plays: int

class TrackPlaysRequest(BaseModel):
    id: str
    plays: int

# Connection Manager for WebSockets
class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)

    async def broadcast(self, message: dict):
        for connection in list(self.active_connections):
            try:
                await connection.send_json(message)
            except Exception:
                self.disconnect(connection)

manager = ConnectionManager()

# Helper to determine active guild
def get_active_guild():
    voice_client = bot.voice_clients[0] if bot.voice_clients else None
    if voice_client:
        return voice_client.guild
    elif bot.guilds:
        return bot.guilds[0]
    return None

# FastAPI Endpoints
@app.get("/")
async def root():
    """Service descriptor - the panel is served by its own container."""
    return {
        "service": "zyrarbot-api",
        "docs": "/docs",
        "status": "/api/status",
        "websocket": "/ws",
    }

@app.get("/api/health")
async def health():
    return {"status": "ok", "discord_ready": bot.is_ready()}

@app.get("/api/status")
async def get_status():
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        return {"error": "MusicBot cog not loaded"}
        
    all_guilds = []
    for g in bot.guilds:
        g_tc = []
        for tc in g.text_channels:
            if tc.permissions_for(g.me).send_messages:
                g_tc.append({
                    "channel_id": str(tc.id),
                    "channel_name": tc.name
                })
        active_notif = music_cog.get_notification_channel(g.id)
        all_guilds.append({
            "guild_id": str(g.id),
            "guild_name": g.name,
            "text_channels": g_tc,
            "active_notification_channel_id": str(active_notif.id) if active_notif else None
        })

    all_voice_channels = []
    for g in bot.guilds:
        for vc in g.voice_channels:
            all_voice_channels.append({
                "guild_name": g.name,
                "guild_id": str(g.id),
                "channel_name": vc.name,
                "channel_id": str(vc.id)
            })

    guild = get_active_guild()
    if not guild:
        return {
            "playing": False,
            "paused": False,
            "loop_single": False,
            "loop_playlist": False,
            "volume": 100,
            "current": None,
            "queue": [],
            "history": [],
            "voice_channels": all_voice_channels,
            "active_voice_channel": None,
            "text_channels": [],
            "active_text_channel_id": None,
            "guilds": all_guilds,
            "default_plays": music_cog.get_default_plays() if music_cog else 5
        }
        
    guild_id = guild.id
    
    # Active voice channel
    voice_client = discord.utils.get(bot.voice_clients, guild__id=guild_id)
    active_vc = None
    if voice_client and voice_client.channel:
        active_vc = {
            "channel_name": voice_client.channel.name,
            "channel_id": str(voice_client.channel.id),
            "guild_name": voice_client.channel.guild.name,
            "guild_id": str(voice_client.channel.guild.id)
        }
        
    # Gather text channels available to the bot on this active guild where it can write
    text_channels = []
    for channel in guild.text_channels:
        if channel.permissions_for(guild.me).send_messages:
            text_channels.append({
                "guild_name": guild.name,
                "guild_id": str(guild.id),
                "channel_name": channel.name,
                "channel_id": str(channel.id)
            })
                
    active_tc_id = None
    active_tc = music_cog.get_notification_channel(guild_id)
    if active_tc:
        active_tc_id = str(active_tc.id)
        
    playing = voice_client.is_playing() if voice_client else False
    paused = voice_client.is_paused() if voice_client else False
    
    current_track = music_cog.get_current(guild_id)
    queue = music_cog.get_queue(guild_id)
    
    return {
        "playing": playing,
        "paused": paused,
        "loop_single": music_cog.get_loop_current(guild_id),
        "loop_playlist": music_cog.get_loop_playlist(guild_id),
        "volume": int(music_cog.get_volume(guild_id) * 100),
        "current": {
            "title": current_track.title,
            "url": current_track.data.get("webpage_url") or current_track.url,
            "id": current_track.id,
            "duration": current_track.duration,
            "elapsed": round(current_track.elapsed, 1),
            "plays_left": getattr(current_track, 'plays_left', music_cog.get_default_plays())
        } if current_track else None,
        "queue": [
            {
                "title": item.title,
                "url": item.data.get("webpage_url") or item.url,
                "id": item.id,
                "plays_left": getattr(item, 'plays_left', music_cog.get_default_plays())
            } for item in queue
        ],
        "queue_index": music_cog.get_queue_index(guild_id),
        "default_plays": music_cog.get_default_plays(),
        "voice_channels": all_voice_channels,
        "active_voice_channel": active_vc,
        "text_channels": text_channels,
        "active_text_channel_id": active_tc_id,
        "guilds": all_guilds,
        "history": [
            {
                "id": item["id"],
                "title": item["title"],
                "url": item["url"],
                "duration": item.get("duration"),
                "is_local": item.get("is_local", False),
                "filepath": item.get("filepath")
            } for item in music_cog.load_history(guild_id)[:20]
        ]
    }

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        # Immediately send status upon connection
        status = await get_status()
        await websocket.send_json(status)
        
        while True:
            # Maintain connection, listen for disconnect
            await websocket.receive_text()
    except WebSocketDisconnect:
        manager.disconnect(websocket)
    except Exception as e:
        print(f"WS connection error: {e}")
        manager.disconnect(websocket)

async def broadcast_status():
    if not manager.active_connections:
        return
    try:
        status = await get_status()
        await manager.broadcast(status)
    except Exception as e:
        print(f"Error in broadcast_status: {e}")

@app.post("/api/play")
async def api_play(req: PlayRequest):
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
        
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    guild_id = guild.id
    voice_client = discord.utils.get(bot.voice_clients, guild__id=guild_id)
    if not voice_client:
        raise HTTPException(status_code=400, detail="Bot is not connected to a voice channel on this server")
        
    try:
        filepath = os.path.join("uploads", os.path.basename(req.query))
        if os.path.exists(filepath) and not req.query.startswith("http"):
            player = await YTDLSource.from_file(filepath, loop=bot.loop)
        else:
            player = await YTDLSource.from_url(req.query, loop=bot.loop, stream=True)
        
        res = await music_cog.enqueue_track(guild_id, player, from_web=True)
        await broadcast_status()
        return res
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/pause")
async def api_pause():
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    voice_client = discord.utils.get(bot.voice_clients, guild__id=guild.id)
    if voice_client and voice_client.is_playing():
        voice_client.pause()
        await broadcast_status()
        return {"status": "paused"}
    raise HTTPException(status_code=400, detail="Bot is not playing on this server")

@app.post("/api/resume")
async def api_resume():
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    voice_client = discord.utils.get(bot.voice_clients, guild__id=guild.id)
    if voice_client and voice_client.is_paused():
        voice_client.resume()
        await broadcast_status()
        return {"status": "resumed"}
    raise HTTPException(status_code=400, detail="Bot is not paused on this server")

@app.post("/api/skip")
async def api_skip():
    music_cog = bot.get_cog("MusicBot")
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    if music_cog and music_cog.get_loop_current(guild.id):
        music_cog.set_loop_current(guild.id, False)

    voice_client = discord.utils.get(bot.voice_clients, guild__id=guild.id)
    if voice_client and (voice_client.is_playing() or voice_client.is_paused()):
        voice_client.stop()
        await broadcast_status()
        return {"status": "skipped"}
    raise HTTPException(status_code=400, detail="Bot is not playing anything on this server")

@app.post("/api/jump")
async def api_jump(req: JumpRequest):
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
    guild_id = guild.id
    queue = music_cog.get_queue(guild_id)
    if req.index < 0 or req.index >= len(queue):
        raise HTTPException(status_code=400, detail="Invalid track index")
        
    if music_cog.get_loop_current(guild_id):
        music_cog.set_loop_current(guild_id, False)
        
    music_cog.set_queue_index(guild_id, req.index)
    music_cog.jump_flag[guild_id] = True
    
    voice_client = discord.utils.get(bot.voice_clients, guild__id=guild_id)
    if voice_client and (voice_client.is_playing() or voice_client.is_paused()):
        voice_client.stop()
    else:
        await music_cog.play_next(guild_id, advance=False)
        
    await broadcast_status()
    return {"status": "jumped"}

@app.post("/api/stop")
async def api_stop():
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
        
    guild = get_active_guild()
    if not guild:
        return {"status": "stopped"}
        
    guild_id = guild.id
    voice_client = discord.utils.get(bot.voice_clients, guild__id=guild_id)
    if voice_client and (voice_client.is_playing() or voice_client.is_paused()):
        voice_client.stop()
        
    music_cog.clear_queue(guild_id)
    await broadcast_status()
    return {"status": "stopped"}

@app.post("/api/volume")
async def api_volume(req: VolumeRequest):
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
        
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    guild_id = guild.id
    if 0 <= req.volume <= 100:
        music_cog.set_volume(guild_id, req.volume / 100)
        voice_client = discord.utils.get(bot.voice_clients, guild__id=guild_id)
        if voice_client and voice_client.source:
            voice_client.source.volume = music_cog.get_volume(guild_id)
        await broadcast_status()
        return {"status": "success", "volume": req.volume}
    raise HTTPException(status_code=400, detail="Volume must be between 0 and 100")

@app.post("/api/loop")
async def api_loop(req: LoopRequest):
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
        
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    music_cog.set_loop_current(guild.id, req.loop)
    if req.loop:
        music_cog.set_loop_playlist(guild.id, False)
    await broadcast_status()
    return {"status": "success", "loop": req.loop}

@app.post("/api/loop_playlist")
async def api_loop_playlist():
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
        
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    current_state = music_cog.get_loop_playlist(guild.id)
    new_state = not current_state
    music_cog.set_loop_playlist(guild.id, new_state)
    if new_state:
        music_cog.set_loop_current(guild.id, False)
    await broadcast_status()
    return {"status": "success", "loop_playlist": new_state}

@app.post("/api/loop_cycle")
async def api_loop_cycle():
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
        
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    loop_current = music_cog.get_loop_current(guild.id)
    loop_pl = music_cog.get_loop_playlist(guild.id)
    
    if not loop_current and not loop_pl:
        music_cog.set_loop_playlist(guild.id, True)
        music_cog.set_loop_current(guild.id, False)
        mode = "playlist"
    elif loop_pl and not loop_current:
        music_cog.set_loop_playlist(guild.id, False)
        music_cog.set_loop_current(guild.id, True)
        mode = "single"
    else:
        music_cog.set_loop_playlist(guild.id, False)
        music_cog.set_loop_current(guild.id, False)
        mode = "off"
        
    await broadcast_status()
    return {"status": "success", "mode": mode}

@app.post("/api/remove")
async def api_remove(req: RemoveRequest):
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
        
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    guild_id = guild.id
    queue = music_cog.get_queue(guild_id)
    current_idx = music_cog.get_queue_index(guild_id)
    
    removed_idx = -1
    for idx, item in enumerate(queue):
        if item.id == req.id:
            removed_idx = idx
            break
            
    if removed_idx == -1:
        raise HTTPException(status_code=404, detail="Track not found in queue")
        
    if removed_idx == current_idx:
        new_queue = [item for idx, item in enumerate(queue) if idx != removed_idx]
        music_cog.queues[guild_id] = new_queue
        music_cog.set_queue_index(guild_id, current_idx)
        music_cog.jump_flag[guild_id] = True
        voice_client = discord.utils.get(bot.voice_clients, guild__id=guild_id)
        if voice_client and (voice_client.is_playing() or voice_client.is_paused()):
            voice_client.stop()
        else:
            await music_cog.play_next(guild_id, advance=False)
    else:
        new_queue = [item for idx, item in enumerate(queue) if idx != removed_idx]
        music_cog.queues[guild_id] = new_queue
        if removed_idx < current_idx:
            music_cog.set_queue_index(guild_id, current_idx - 1)
            
    await broadcast_status()
    return {"status": "removed"}

@app.post("/api/reorder")
async def api_reorder(req: ReorderRequest):
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
        
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    guild_id = guild.id
    queue = music_cog.get_queue(guild_id)
    current_idx = music_cog.get_queue_index(guild_id)
    curr_item = queue[current_idx] if 0 <= current_idx < len(queue) else None
    
    queue_map = {item.id: item for item in queue}
    
    new_queue = []
    for track_id in req.queue_ids:
        if track_id in queue_map:
            new_queue.append(queue_map[track_id])
            
    # Retain any items not present in the reordered IDs list
    for item in queue:
        if item.id not in req.queue_ids:
            new_queue.append(item)
            
    music_cog.queues[guild_id] = new_queue
    if curr_item:
        try:
            music_cog.set_queue_index(guild_id, new_queue.index(curr_item))
        except ValueError:
            pass
            
    await broadcast_status()
    return {"status": "success"}

@app.get("/api/uploads")
async def get_uploads():
    uploads_dir = "uploads"
    os.makedirs(uploads_dir, exist_ok=True)
    files = []
    for f in os.listdir(uploads_dir):
        filepath = os.path.join(uploads_dir, f)
        if os.path.isfile(filepath) and not f.startswith("."):
            size_mb = round(os.path.getsize(filepath) / (1024 * 1024), 2)
            mtime = os.path.getmtime(filepath)
            files.append({
                "filename": f,
                "size_mb": size_mb,
                "mtime": mtime
            })
    files.sort(key=lambda x: x["mtime"], reverse=True)
    return {"files": files}

@app.post("/api/upload")
async def api_upload_file(file: UploadFile = File(...)):
    uploads_dir = "uploads"
    os.makedirs(uploads_dir, exist_ok=True)
    
    filename = os.path.basename(file.filename)
    filepath = os.path.join(uploads_dir, filename)
    
    try:
        with open(filepath, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        return {"status": "success", "filename": filename}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Błąd zapisu pliku: {e}")

@app.delete("/api/uploads/{filename}")
async def api_delete_file(filename: str):
    uploads_dir = "uploads"
    filepath = os.path.join(uploads_dir, os.path.basename(filename))
    if os.path.exists(filepath):
        try:
            os.remove(filepath)
            return {"status": "deleted"}
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))
    raise HTTPException(status_code=404, detail="Plik nie istnieje")

@app.post("/api/play_file")
async def api_play_file(req: PlayFileRequest):
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
        
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    guild_id = guild.id
    voice_client = discord.utils.get(bot.voice_clients, guild__id=guild_id)
    if not voice_client:
        raise HTTPException(status_code=400, detail="Bot is not connected to a voice channel on this server")
        
    filepath = os.path.join("uploads", os.path.basename(req.filename))
    if not os.path.exists(filepath):
        raise HTTPException(status_code=404, detail="Plik nie istnieje na serwerze")
        
    try:
        os.utime(filepath, None)
        player = await YTDLSource.from_file(filepath, loop=bot.loop)
        res = await music_cog.enqueue_track(guild_id, player, from_web=True)
        await broadcast_status()
        return res
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/connect")
async def api_connect(req: ConnectRequest):
    try:
        channel_id = int(req.channel_id)
    except ValueError:
        raise HTTPException(status_code=400, detail="Invalid channel ID")
        
    channel = bot.get_channel(channel_id)
    if not channel:
        raise HTTPException(status_code=404, detail="Voice channel not found")
        
    guild_id = channel.guild.id
    
    # Only disconnect voice clients from the SAME guild!
    for vc in bot.voice_clients:
        if vc.guild.id == guild_id:
            await vc.disconnect()
            
    await channel.connect()
    
    music_cog = bot.get_cog("MusicBot")
    if music_cog and not music_cog.get_notification_channel(guild_id):
        # Save first accessible text channel to send updates
        for tc in channel.guild.text_channels:
            if tc.permissions_for(channel.guild.me).send_messages:
                music_cog.last_channels[guild_id] = tc
                break
                
    await broadcast_status()
    return {"status": "connected", "channel": channel.name}

@app.post("/api/disconnect")
async def api_disconnect():
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
        
    guild_id = guild.id
    voice_client = discord.utils.get(bot.voice_clients, guild__id=guild_id)
    if voice_client:
        await voice_client.disconnect()
            
        music_cog = bot.get_cog("MusicBot")
        if music_cog:
            music_cog.clear_queue(guild_id)
            
        await broadcast_status()
        return {"status": "disconnected"}
    raise HTTPException(status_code=400, detail="Bot not connected to voice on this server")

@app.post("/api/settings/text_channel")
async def api_set_text_channel(req: TextChannelRequest):
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
        
    if not req.channel_id or req.channel_id == "0" or req.channel_id == "":
        if not req.guild_id:
            raise HTTPException(status_code=400, detail="guild_id required to reset channel")
        guild_id = int(req.guild_id)
        if guild_id in music_cog.last_channels:
            del music_cog.last_channels[guild_id]
        music_cog.save_channel_id(guild_id, None)
        await broadcast_status()
        return {"status": "success", "channel_id": None, "guild_id": str(guild_id)}

    channel = bot.get_channel(int(req.channel_id))
    if not channel:
        raise HTTPException(status_code=404, detail="Text channel not found")
        
    guild_id = channel.guild.id
    if req.guild_id and int(req.guild_id) != guild_id:
        raise HTTPException(status_code=400, detail="Channel does not belong to the specified guild")
        
    music_cog.last_channels[guild_id] = channel
    music_cog.save_channel_id(guild_id, int(req.channel_id))
    
    # Send test message to confirm
    try:
        await channel.send(f"🔔 Ten kanał został ustawiony jako główny kanał powiadomień dla ŻyrarBOT na serwerze **{channel.guild.name}** z panelu webowego.")
    except Exception as e:
        print(f"Failed to send test message: {e}")
        
    await broadcast_status()
    return {"status": "success", "channel_id": req.channel_id, "guild_id": str(guild_id)}

@app.post("/api/settings/default_plays")
async def api_set_default_plays(req: DefaultPlaysRequest):
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
    if req.default_plays < 1 or req.default_plays > 999:
        raise HTTPException(status_code=400, detail="Invalid default plays value")
    music_cog.set_default_plays(req.default_plays)
    await broadcast_status()
    return {"status": "success", "default_plays": req.default_plays}

@app.post("/api/queue/plays")
async def api_set_track_plays(req: TrackPlaysRequest):
    music_cog = bot.get_cog("MusicBot")
    if not music_cog:
        raise HTTPException(status_code=500, detail="MusicBot cog not loaded")
    guild = get_active_guild()
    if not guild:
        raise HTTPException(status_code=400, detail="Bot has no active server context")
    guild_id = guild.id
    queue = music_cog.get_queue(guild_id)
    if req.plays < 1 or req.plays > 999:
        raise HTTPException(status_code=400, detail="Invalid plays value")
    found = False
    for item in queue:
        if item.id == req.id:
            item.plays_left = req.plays
            found = True
            break
    if not found:
        raise HTTPException(status_code=404, detail="Track not found in queue")
    await broadcast_status()
    return {"status": "success", "id": req.id, "plays": req.plays}

# WebSocket Background Loop and Startup Hook
async def status_broadcast_loop():
    while True:
        try:
            await asyncio.sleep(1)
            if manager.active_connections:
                await broadcast_status()
        except asyncio.CancelledError:
            break
        except Exception as e:
            print(f"Error in status_broadcast_loop: {e}")

async def wait_and_register_callback():
    for _ in range(30):  # Try for 30 seconds
        await asyncio.sleep(1)
        music_cog = bot.get_cog("MusicBot")
        if music_cog:
            music_cog.on_state_change_callback = broadcast_status
            print("Registered WebSocket state change callback on MusicBot cog (delayed)")
            break

@app.on_event("startup")
async def startup_event():
    music_cog = bot.get_cog("MusicBot")
    if music_cog:
        music_cog.on_state_change_callback = broadcast_status
        print("Registered WebSocket state change callback on MusicBot cog")
    else:
        asyncio.create_task(wait_and_register_callback())
    
    asyncio.create_task(status_broadcast_loop())