backend/main.py

import os
import sys
import io
import socket
import subprocess
import threading
import mimetypes
import json
from pathlib import Path
from datetime import datetime
from typing import Optional

sys.stdout = open(sys.stdout.fileno(), mode='w', encoding='utf-8', buffering=1)

import qrcode
import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse, JSONResponse, Response, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from dotenv import load_dotenv
from pyngrok import ngrok, conf

load_dotenv()

# ─── Config ───────────────────────────────────────────────────────────────────
PORT         = int(os.environ.get("PORT", 8000))
BASE_DIR     = Path(__file__).resolve().parent
FRONTEND_DIR = BASE_DIR.parent / "frontend"
UPLOAD_DIR   = BASE_DIR / "uploads"
URL_FILE     = BASE_DIR / "public_url.txt"
UPLOAD_DIR.mkdir(exist_ok=True)

# ─── Rooms State ─────────────────────────────────────────────────────────────
# { "room_code": { "clients": [WebSocket], "messages": [dict] } }
rooms = {}

# ─── Kill any process already on PORT ────────────────────────────────────────
def free_port(port: int):
    try:
        result = subprocess.run(["netstat", "-ano"], capture_output=True, text=True)
        for line in result.stdout.splitlines():
            if f":{port}" in line and "LISTENING" in line:
                pid = line.strip().split()[-1]
                subprocess.run(["taskkill", "/F", "/PID", pid], capture_output=True)
                print(f"  [~] Freed port {port} (killed PID {pid})")
    except Exception:
        pass

# ─── App ─────────────────────────────────────────────────────────────────────
app = FastAPI(title="RadarDrop v2")
app.add_middleware(CORSMiddleware, allow_origins=["*"],
                   allow_methods=["*"], allow_headers=["*"])

# ─── Helpers ─────────────────────────────────────────────────────────────────
def get_local_ip() -> str:
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        return "127.0.0.1"

def human_size(n: int) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024:
            return f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} TB"

def make_qr_png(url: str) -> bytes:
    img = qrcode.make(url)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    return buf.read()

def get_public_url() -> Optional[str]:
    try:
        txt = URL_FILE.read_text().strip()
        return txt if txt else None
    except Exception:
        return None

# ─── Ngrok Tunnel ────────────────────────────────────────────────────────────
def start_ngrok_tunnel(port: int):
    URL_FILE.write_text("", encoding="utf-8")
    
    auth_token = os.environ.get("NGROK_AUTH_TOKEN")
    if not auth_token:
        print("\n  [!] NGROK_AUTH_TOKEN is missing in .env file.")
        print("  [!] Running strictly on Local Network.")
        URL_FILE.write_text("FAILED", encoding="utf-8")
        return

    try:
        ngrok.set_auth_token(auth_token)
        # Suppress verbose ngrok logs
        conf.get_default().log_event_callback = lambda log: None
        
        domain = os.environ.get("NGROK_DOMAIN")
        if domain:
            tunnel = ngrok.connect(port, domain=domain)
        else:
            tunnel = ngrok.connect(port)
            
        public_url = tunnel.public_url
        URL_FILE.write_text(public_url, encoding="utf-8")
        print(f"\n  [PUBLIC] Tunnel live: {public_url}")
        print(f"  [PUBLIC] Share with anyone!\n")
    except Exception as e:
        print(f"\n  [WARN] Ngrok failed: {e}")
        URL_FILE.write_text("FAILED", encoding="utf-8")

async def broadcast_to_room(room_code: str, message: dict):
    if room_code in rooms:
        for client in rooms[room_code]["clients"]:
            try:
                await client.send_json(message)
            except Exception:
                pass

# ─── WebSockets (Chat & Live Sync) ───────────────────────────────────────────
@app.websocket("/ws/{room_code}")
async def websocket_endpoint(websocket: WebSocket, room_code: str):
    await websocket.accept()
    if room_code not in rooms:
        rooms[room_code] = {"clients": [], "messages": []}
        (UPLOAD_DIR / room_code).mkdir(exist_ok=True)
    
    rooms[room_code]["clients"].append(websocket)
    
    # Send chat history to new connection
    for msg in rooms[room_code]["messages"]:
        await websocket.send_json(msg)
        
    try:
        while True:
            data = await websocket.receive_text()
            payload = json.loads(data)
            if payload.get("type") == "chat":
                msg = {
                    "type": "chat",
                    "sender": payload.get("sender", "Anonymous"),
                    "text": payload.get("text"),
                    "timestamp": datetime.now().strftime("%H:%M")
                }
                rooms[room_code]["messages"].append(msg)
                await broadcast_to_room(room_code, msg)
    except WebSocketDisconnect:
        rooms[room_code]["clients"].remove(websocket)
    except Exception as e:
        if websocket in rooms[room_code]["clients"]:
            rooms[room_code]["clients"].remove(websocket)

# ─── Routes ──────────────────────────────────────────────────────────────────

@app.api_route("/", methods=["GET", "HEAD"], response_class=HTMLResponse)
async def serve_index():
    return HTMLResponse((FRONTEND_DIR / "index.html").read_text(encoding="utf-8"))

@app.get("/ip")
async def get_ip():
    ip         = get_local_ip()
    public_url = get_public_url()
    failed     = public_url == "FAILED"
    return {
        "local_url":     f"http://{ip}:{PORT}",
        "public_url":    None if (not public_url or failed) else public_url,
        "tunnel_status": "failed" if failed else ("live" if public_url else "connecting"),
    }

@app.get("/qr")
async def local_qr(room: str = ""):
    url = f"http://{get_local_ip()}:{PORT}"
    if room: url += f"?room={room}"
    return Response(content=make_qr_png(url), media_type="image/png")

@app.get("/qr/public")
async def public_qr(room: str = ""):
    url = get_public_url()
    if not url or url == "FAILED":
        raise HTTPException(503, "Tunnel not ready yet")
    if room: url += f"?room={room}"
    return Response(content=make_qr_png(url), media_type="image/png")

# ─── Room File Routes ────────────────────────────────────────────────────────

@app.get("/api/rooms/{room_code}/files")
async def list_files(room_code: str):
    room_dir = UPLOAD_DIR / room_code
    room_dir.mkdir(exist_ok=True)
    
    files = []
    for f in sorted(room_dir.iterdir(), key=lambda x: x.stat().st_mtime, reverse=True):
        if f.is_file():
            stat = f.stat()
            mime, _ = mimetypes.guess_type(f.name)
            files.append({
                "name":     f.name,
                "size":     stat.st_size,
                "size_hr":  human_size(stat.st_size),
                "mime":     mime or "application/octet-stream",
                "modified": datetime.fromtimestamp(stat.st_mtime).strftime("%Y-%m-%d %H:%M"),
            })
    return JSONResponse(files)

@app.post("/api/rooms/{room_code}/upload")
async def upload(room_code: str, file: UploadFile = File(...)):
    room_dir = UPLOAD_DIR / room_code
    room_dir.mkdir(exist_ok=True)
    
    dest = room_dir / file.filename
    stem = Path(file.filename).stem
    suffix = Path(file.filename).suffix
    n = 1
    while dest.exists():
        dest = room_dir / f"{stem} ({n}){suffix}"; n += 1
        
    dest.write_bytes(await file.read())
    
    # Broadcast to room that files updated
    await broadcast_to_room(room_code, {"type": "files_updated"})
    
    return {"filename": dest.name}

@app.get("/api/rooms/{room_code}/download/{filename}")
async def download(room_code: str, filename: str):
    p = UPLOAD_DIR / room_code / filename
    if not p.exists():
        raise HTTPException(404, "Not found")
    return FileResponse(p, filename=filename)

@app.delete("/api/rooms/{room_code}/files/{filename}")
async def delete(room_code: str, filename: str):
    p = UPLOAD_DIR / room_code / filename
    if not p.exists():
        raise HTTPException(404, "Not found")
    p.unlink()
    
    await broadcast_to_room(room_code, {"type": "files_updated"})
    return {"ok": True}

# ─── Entry Point ─────────────────────────────────────────────────────────────
if __name__ == "__main__":
    ip = get_local_ip()
    
    # If we are NOT running on Hugging Face Spaces, free port and prompt for Ngrok
    is_huggingface = os.environ.get("SPACE_ID") is not None
    
    if not is_huggingface:
        free_port(PORT)
        if not os.environ.get("NGROK_AUTH_TOKEN"):
        print("\n" + "="*55)
        print("  [!] FIXED URL SETUP REQUIRED")
        print("  To get a fixed URL, you need a free Ngrok account.")
        print("  1. Go to: https://dashboard.ngrok.com/get-started/your-authtoken")
        print("  2. Copy your Authtoken.")
        print("  3. (Optional) Claim your free static domain at: https://dashboard.ngrok.com/cloud-edge/domains")
        print("  (Or leave blank and press Enter to use Local Network only)")
        print("="*55)
        token = input("  Enter Ngrok Authtoken: ").strip()
        if token:
            domain = input("  (Optional) Enter your Ngrok Static Domain (e.g., xxx.ngrok-free.app) or leave blank: ").strip()
            with open(BASE_DIR / ".env", "w") as f:
                f.write(f"NGROK_AUTH_TOKEN={token}\n")
                if domain:
                    f.write(f"NGROK_DOMAIN={domain}\n")
            os.environ["NGROK_AUTH_TOKEN"] = token
            if domain:
                os.environ["NGROK_DOMAIN"] = domain
            print("  [+] Saved to .env!")
        else:
            print("  [-] No token provided. Running in Local Mode only.")

        threading.Thread(target=start_ngrok_tunnel, args=(PORT,), daemon=True).start()

    print("\n" + "=" * 55)
    print("  RadarDrop v2 is LIVE")
    print(f"  Local  : http://localhost:{PORT}")
    if is_huggingface:
        print("  Running in Hugging Face Space")
    else:
        print(f"  Network: http://{ip}:{PORT}")
    print("=" * 55 + "\n")

    uvicorn.run(app, host="0.0.0.0", port=PORT)