from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from typing import Dict
import sqlite3
import json
import os

app = FastAPI()

# Enable CORS so your mobile app/web app can connect without security blocks
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Your mounted storage bucket path
DB_PATH = "/data/lowkey.db"

def init_db():
    # Ensures the database file is created directly inside your 1GB persistent bucket
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS queue (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            sender TEXT,
            receiver TEXT,
            message TEXT,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()
    conn.close()

# Initialize DB on server startup
@app.on_event("startup")
async def startup_event():
    init_db()

class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}

    async def connect(self, user_id: str, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[user_id] = websocket
        await self.flush_offline_messages(user_id)

    def disconnect(self, user_id: str):
        if user_id in self.active_connections:
            del self.active_connections[user_id]

    async def send_message(self, receiver_id: str, message: dict):
        if receiver_id in self.active_connections:
            try:
                await self.active_connections[receiver_id].send_text(json.dumps(message))
                return True
            except Exception:
                # If sending fails halfway through, fallback to saving offline
                self.queue_message(message['sender'], receiver_id, message['text'])
                return False
        else:
            self.queue_message(message['sender'], receiver_id, message['text'])
            return False

    def queue_message(self, sender, receiver, text):
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        cursor.execute("INSERT INTO queue (sender, receiver, message) VALUES (?, ?, ?)", (sender, receiver, text))
        conn.commit()
        conn.close()

    async def flush_offline_messages(self, user_id: str):
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        cursor.execute("SELECT id, sender, message FROM queue WHERE receiver = ?", (user_id,))
        rows = cursor.fetchall()
        
        if rows:
            websocket = self.active_connections[user_id]
            for row in rows:
                msg_id, sender, text = row
                payload = {"sender": sender, "text": text, "status": "offline_delivery"}
                try:
                    await websocket.send_text(json.dumps(payload))
                except Exception:
                    conn.close()
                    return # Stop if user disconnects mid-delivery
            
            cursor.execute("DELETE FROM queue WHERE receiver = ?", (user_id,))
            conn.commit()
        conn.close()

manager = ConnectionManager()

@app.get("/")
def read_root():
    return {"status": "Lowkey Server is up and running 🔥"}

@app.websocket("/ws/{user_id}")
async def websocket_endpoint(websocket: WebSocket, user_id: str):
    await manager.connect(user_id, websocket)
    try:
        while True:
            data = await websocket.receive_text()
            payload = json.loads(data)
            
            # Action: App client says "I downloaded and saved this safely"
            if payload.get("action") == "seen_and_saved":
                continue
            
            target_receiver = payload.get("receiver")
            msg_data = {
                "sender": user_id,
                "text": payload.get("text")
            }
            await manager.send_message(target_receiver, msg_data)
            
    except WebSocketDisconnect:
        manager.disconnect(user_id)