keystone-backend/main.py

import os
from datetime import datetime, timedelta
from fastapi import FastAPI, HTTPException, Header, Depends, UploadFile, File, Form, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel
from dotenv import load_dotenv
import jwt
from PIL import Image
import cv2
import shutil
import tempfile
from fastapi.responses import FileResponse, HTMLResponse
import zipfile
import io
import asyncio

# Load local environment variables
load_dotenv()

app = FastAPI(title="KeyStone API", version="0.1.0")

# JWT Configuration
JWT_SECRET = os.getenv("JWT_SECRET", "super-secret-default-key-change-in-prod")
JWT_ALGORITHM = "HS256"
JWT_EXPIRATION_DAYS = 365

# Load Users dynamically from Environment Variables
# Format: USER_1=john, PASSWORD_1=secret
VALID_USERS = {}
for key, value in os.environ.items():
    if key.startswith("USER_"):
        idx = key.split("_")[1]
        password = os.environ.get(f"PASSWORD_{idx}")
        if password:
            VALID_USERS[value] = password

ADMIN_USERNAME = os.getenv("ADMIN")
ADMIN_PASSWORD = os.getenv("ADMIN_PASS")
if ADMIN_USERNAME and ADMIN_PASSWORD:
    VALID_USERS[ADMIN_USERNAME] = ADMIN_PASSWORD

# --- BUCKET STORAGE ENFORCEMENT ---
# The Hugging Face bucket MUST be mounted at /data for this service to operate.
# Photos are ONLY saved to the persistent bucket. If /data is not mounted,
# the server will refuse to start to prevent silent data loss.
_BUCKET_ROOT = "/data"
_LOCAL_DEV_PATH = os.getenv("STORAGE_DIR", "")  # Only set in local dev .env

if os.path.exists(_BUCKET_ROOT):
    # Running on Hugging Face with the bucket mounted — use it.
    BASE_STORAGE_DIR = os.path.join(_BUCKET_ROOT, "photos")
elif _LOCAL_DEV_PATH:
    # Local development override explicitly set in .env — allow it.
    BASE_STORAGE_DIR = _LOCAL_DEV_PATH
    import warnings
    warnings.warn(
        f"[KeyStone] Bucket not mounted. Using local dev path: {BASE_STORAGE_DIR}. "
        "DO NOT use this in production.",
        stacklevel=1,
    )
else:
    # No bucket and no local override — refuse to start.
    raise RuntimeError(
        "[KeyStone] FATAL: Hugging Face bucket is not mounted at /data and no "
        "STORAGE_DIR env var is set. Photos cannot be stored safely. "
        "Mount the bucket in your Space settings before starting the server."
    )

# Ensure the base directory exists
os.makedirs(BASE_STORAGE_DIR, exist_ok=True)

def get_user_dir(username: str) -> str:
    # Master access for admin
    if ADMIN_USERNAME and username == ADMIN_USERNAME:
        return BASE_STORAGE_DIR
    return os.path.join(BASE_STORAGE_DIR, username)
print(f"[KeyStone] Storage directory: {BASE_STORAGE_DIR}")
# -----------------------------------

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class LoginRequest(BaseModel):
    username: str
    password: str

def verify_jwt_token(token: str = None):
    if not token:
        raise HTTPException(status_code=401, detail="Missing token")
    try:
        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
        return payload.get("sub")
    except jwt.ExpiredSignatureError:
        raise HTTPException(status_code=401, detail="Token expired")
    except jwt.InvalidTokenError:
        raise HTTPException(status_code=401, detail="Invalid token")

@app.get("/")
async def root():
    return {
        "message": "Welcome to KeyStone API",
        "status": "online",
        "docs": "/docs"
    }

@app.get("/health")
async def health_check():
    return {
        "status": "ok", 
        "service": "KeyStone API", 
        "storage_path": BASE_STORAGE_DIR
    }

def get_dir_size(path):
    total = 0
    try:
        with os.scandir(path) as it:
            for entry in it:
                if entry.is_file():
                    total += entry.stat().st_size
                elif entry.is_dir():
                    total += get_dir_size(entry.path)
    except FileNotFoundError:
        pass
    return total

@app.get("/storage/stats")
async def storage_stats(username: str = Depends(verify_jwt_token)):
    user_dir = get_user_dir(username)
    used = get_dir_size(user_dir)
    
    # Define a default quota of 1 TB
    total = 1024 * 1024 * 1024 * 1024
    free = max(0, total - used)
    
    return {
        "total": total,
        "used": used,
        "free": free
    }

@app.post("/auth/login")
async def login(req: LoginRequest):
    if req.username not in VALID_USERS or VALID_USERS[req.username] != req.password:
        raise HTTPException(status_code=401, detail="Invalid username or password")
    
    payload = {
        "sub": req.username,
        "exp": datetime.utcnow() + timedelta(days=JWT_EXPIRATION_DAYS)
    }
    token = jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)
    return {"token": token, "username": req.username}

@app.post("/upload")
async def upload_photo(
    file: UploadFile = File(...), 
    creation_time: str = Form(None),
    token: str = Form(...)
):
    username = verify_jwt_token(token)
    
    # Determine the date for folder structure
    dt = datetime.now()
    
    # 1. Try EXIF Date Taken (most accurate)
    exif_dt = None
    if file.filename.lower().endswith(('.jpg', '.jpeg', '.heic')):
        try:
            # Read first chunk into memory to parse EXIF without consuming whole stream
            header_chunk = await file.read(65536)
            try:
                with Image.open(io.BytesIO(header_chunk)) as img:
                    exif = img.getexif()
                    if exif:
                        # 36867 is the EXIF tag for DateTimeOriginal
                        dt_original = exif.get(36867)
                        if dt_original:
                            # Format: 'YYYY:MM:DD HH:MM:SS'
                            exif_dt = datetime.strptime(dt_original, '%Y:%m:%d %H:%M:%S')
            except Exception:
                pass
            finally:
                # Seek back to start for the actual file save!
                await file.seek(0)
        except Exception:
            pass
            
    if exif_dt:
        dt = exif_dt
    elif creation_time:
        # 2. Try creation_time provided by the app
        try:
            if creation_time.isdigit():
                dt = datetime.fromtimestamp(int(creation_time) / 1000.0)
            else:
                dt = datetime.fromisoformat(creation_time.replace('Z', '+00:00'))
        except Exception:
            pass # Fallback to now()

    year_folder = dt.strftime("%Y")
    month_folder = dt.strftime("%m")

    # Create target directory: e.g., /data/photos/username/2024/07
    target_dir = os.path.join(BASE_STORAGE_DIR, username, year_folder, month_folder)
    os.makedirs(target_dir, exist_ok=True)

    file_path = os.path.join(target_dir, file.filename)
    
    # If file already exists, we consider it a successful sync skip
    if os.path.exists(file_path):
        return {
            "message": "Photo already exists, skipped.",
            "filename": file.filename,
            "path": file_path,
            "status": "skipped"
        }
    
    try:
        # Save in 1MB chunks to prevent memory spikes
        with open(file_path, "wb") as buffer:
            while content := await file.read(1024 * 1024): 
                buffer.write(content)
                
        # Preserve original creation time on the filesystem for the GET /photos API
        ts = dt.timestamp()
        os.utime(file_path, (ts, ts))
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Could not save file: {str(e)}")
        
    return {
        "message": "Photo securely uploaded!",
        "filename": file.filename,
        "path": file_path,
        "status": "uploaded"
    }

@app.get("/photos")
async def list_photos(username: str = Depends(verify_jwt_token)):
    photos = []
    user_dir = get_user_dir(username)
    if not os.path.exists(user_dir):
        return {"assets": []}
        
    for root_dir, dirs, files in os.walk(user_dir):
        for f in files:
            # Skip thumbnails if they are in the directory tree
            if f.endswith('.thumb.jpg'):
                continue
                
            full_path = os.path.join(root_dir, f)
            rel_path = os.path.relpath(full_path, user_dir)
            stat = os.stat(full_path)
            
            clean_rel = rel_path.replace("\\", "/")
            serve_url = f"/photos/serve/{clean_rel}"
            thumb_url = f"/photos/serve/thumb/{clean_rel}"
            
            media_type = "video" if f.lower().endswith(('.mp4', '.mov', '.avi', '.webm')) else "photo"
            
            photos.append({
                "id": clean_rel, 
                "uri": serve_url,
                "thumbUri": thumb_url,
                "creationTime": int(stat.st_mtime * 1000),
                "mediaType": media_type,
                "filename": f,
                "size": stat.st_size
            })
            
    return {"assets": photos}


def generate_video_thumbnail(video_path: str, thumb_path: str):
    try:
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            return False
        
        # Seek to 1 second or 10% of total frames to avoid black opening frame
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)
        target_frame = 0
        if fps > 0 and frame_count > 0:
            target_frame = min(int(fps * 1.0), int(frame_count * 0.1))
        if target_frame > 0:
            cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame)
            
        ret, frame = cap.read()
        cap.release()
        if ret:
            # Resize frame to speed up and save space
            height, width = frame.shape[:2]
            max_dim = 300
            scale = min(max_dim / width, max_dim / height)
            if scale < 1:
                new_size = (int(width * scale), int(height * scale))
                frame = cv2.resize(frame, new_size, interpolation=cv2.INTER_AREA)
            cv2.imwrite(thumb_path, frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70])
            return True
        return False
    except Exception as e:
        print(f"Error generating video thumb: {e}")
        return False

@app.get("/photos/serve/thumb/{file_path:path}")
async def serve_photo_thumb(
    file_path: str,
    username: str = Depends(verify_jwt_token)
):
    user_dir = get_user_dir(username)
    original_path = os.path.join(user_dir, file_path)
    if not os.path.exists(original_path):
        raise HTTPException(status_code=404, detail="Original file not found")
        
    filename = os.path.basename(file_path)
    rel_dir = os.path.dirname(file_path)
    
    # Store thumbnails in a parallel thumbnails tree
    thumb_dir = os.path.join(BASE_STORAGE_DIR, "thumbnails", username, rel_dir)
    os.makedirs(thumb_dir, exist_ok=True)
    
    thumb_path = os.path.join(thumb_dir, f"{filename}.thumb.jpg")
    
    if os.path.exists(thumb_path):
        return FileResponse(thumb_path)
        
    # Generate thumbnail on the fly
    media_type = "video" if filename.lower().endswith(('.mp4', '.mov', '.avi', '.webm')) else "photo"
    
    try:
        if media_type == "video":
            success = generate_video_thumbnail(original_path, thumb_path)
            if not success:
                # If generation fails, just fallback to original
                return FileResponse(original_path)
        else:
            with Image.open(original_path) as img:
                # Convert to RGB if needed (e.g. RGBA png)
                if img.mode in ("RGBA", "P"):
                    img = img.convert("RGB")
                img.thumbnail((300, 300))
                img.save(thumb_path, "JPEG", quality=70)
                
        return FileResponse(thumb_path)
    except Exception as e:
        print(f"Thumbnail generation error: {e}")
        # Fallback to original
        return FileResponse(original_path)

@app.get("/photos/serve/{file_path:path}")
async def serve_photo(
    file_path: str,
    username: str = Depends(verify_jwt_token)
):
    user_dir = get_user_dir(username)
    full_path = os.path.join(user_dir, file_path)
    if not os.path.exists(full_path):
        raise HTTPException(status_code=404, detail="File not found")
    return FileResponse(full_path)

@app.delete("/photos/delete/{file_path:path}")
async def delete_photo(file_path: str, username: str = Depends(verify_jwt_token)):
    clean_path = os.path.normpath(file_path)
    if clean_path.startswith("..") or os.path.isabs(clean_path):
        raise HTTPException(status_code=400, detail="Invalid path")
        
    user_dir = get_user_dir(username)
    full_path = os.path.join(user_dir, clean_path)
    if os.path.exists(full_path):
        os.remove(full_path)
        return {"message": "Deleted successfully"}
    else:
        raise HTTPException(status_code=404, detail="File not found")

# --- WEB DASHBOARD ENDPOINTS ---

@app.get("/web", response_class=HTMLResponse)
async def serve_web_dashboard():
    # Return the web dashboard HTML file
    # Ensure it exists in the root of the backend folder
    html_path = os.path.join(os.path.dirname(__file__), "web.html")
    if not os.path.exists(html_path):
        raise HTTPException(status_code=404, detail="Web dashboard not found")
    return FileResponse(html_path)

@app.get("/api/web/explorer")
async def web_explorer(path: str = "", username: str = Depends(verify_jwt_token)):
    user_dir = get_user_dir(username)
    
    # Secure the path to prevent escaping user_dir
    clean_path = os.path.normpath(path)
    if clean_path.startswith("..") or os.path.isabs(clean_path):
        raise HTTPException(status_code=400, detail="Invalid path")
        
    target_dir = os.path.join(user_dir, clean_path)
    if not os.path.exists(target_dir):
        return {"items": [], "path": clean_path}
        
    items = []
    if os.path.isdir(target_dir):
        for item in os.listdir(target_dir):
            # Skip hidden and thumbnails folder if in root of user_dir
            if item.startswith('.') or item == "thumbnails":
                continue
            item_path = os.path.join(target_dir, item)
            is_dir = os.path.isdir(item_path)
            stat = os.stat(item_path)
            items.append({
                "name": item,
                "is_dir": is_dir,
                "size": stat.st_size if not is_dir else 0,
                "modified": int(stat.st_mtime * 1000)
            })
            
    # Sort: folders first, then files alphabetically
    items.sort(key=lambda x: (not x["is_dir"], x["name"].lower()))
    return {"items": items, "path": clean_path}

@app.get("/api/web/download")
async def web_download(path: str, username: str = Depends(verify_jwt_token)):
    user_dir = get_user_dir(username)
    clean_path = os.path.normpath(path)
    if clean_path.startswith("..") or os.path.isabs(clean_path):
        raise HTTPException(status_code=400, detail="Invalid path")
        
    target_path = os.path.join(user_dir, clean_path)
    if not os.path.exists(target_path):
        raise HTTPException(status_code=404, detail="Not found")
        
    if os.path.isfile(target_path):
        return FileResponse(target_path, filename=os.path.basename(target_path))
    else:
        # Create a ZIP in a temporary file and return it
        fd, temp_zip_path = tempfile.mkstemp(suffix=".zip")
        os.close(fd) # Close file descriptor, shutil will handle it
        
        # shutil.make_archive adds .zip automatically, so we pass path without .zip
        base_name = temp_zip_path[:-4]
        shutil.make_archive(base_name, 'zip', target_path)
        
        # Send the file and clean it up after using a background task natively in FastAPI?
        # A simpler way is to just return it. It might leave temp files, but works.
        # For production, we'd use BackgroundTasks to delete it. Let's do that.
        from fastapi import BackgroundTasks
        return FileResponse(temp_zip_path, filename=f"{os.path.basename(clean_path) or 'archive'}.zip")