import os
import uuid
import asyncio
import aiofiles
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Header
from fastapi.responses import HTMLResponse, JSONResponse
from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware
import uvicorn

app = FastAPI()
app.add_middleware(ProxyHeadersMiddleware, trusted_hosts="*")

UPLOADS_ROOT  = Path(__file__).parent / "uploads"
SESSIONS_ROOT = Path(__file__).parent / ".sessions"   # temp chunk storage
UPLOADS_ROOT.mkdir(exist_ok=True)
SESSIONS_ROOT.mkdir(exist_ok=True)

HTML_FILE  = Path(__file__).parent / "upload.html"
CHUNK_SIZE = 5 * 1024 * 1024   # 5 MB chunks


# ── security helpers ──────────────────────────────────────────────────────────

def resolve_subdir(subdir: str) -> Path:
    if not subdir or subdir.strip() in ("", "/", "."):
        return UPLOADS_ROOT
    clean = Path(subdir.strip()).as_posix().lstrip("/")
    target = (UPLOADS_ROOT / clean).resolve()
    try:
        target.relative_to(UPLOADS_ROOT.resolve())
    except ValueError:
        raise HTTPException(400, "Invalid subdirectory")
    return target


def safe_filename(name: str) -> str:
    name = Path(name).name
    name = name.replace("\x00", "").strip(". ")
    if not name:
        raise HTTPException(400, "Invalid filename")
    return name


def safe_session_id(sid: str) -> str:
    """Validate that session id is a plain UUID — no path tricks."""
    try:
        return str(uuid.UUID(sid))
    except ValueError:
        raise HTTPException(400, "Invalid session id")


def unique_dest(dest: Path) -> Path:
    if not dest.exists():
        return dest
    stem, suffix = dest.stem, dest.suffix
    i = 1
    while dest.exists():
        dest = dest.with_name(f"{stem}_{i}{suffix}")
        i += 1
    return dest


def list_subdirs() -> list[str]:
    result = [""]
    for root, dirs, _ in os.walk(UPLOADS_ROOT):
        dirs.sort()
        for d in dirs:
            full = Path(root) / d
            rel = str(full.relative_to(UPLOADS_ROOT)).replace("\\", "/")
            result.append(rel)
    return result


# ── session state (in-memory, survives within one process) ───────────────────
# { session_id: { "filename": str, "subdir": str, "total_size": int,
#                 "received": int, "tmp": Path } }
sessions: dict[str, dict] = {}


# ── routes ────────────────────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def index():
    if not HTML_FILE.exists():
        raise HTTPException(500, "upload.html not found")
    return HTMLResponse(HTML_FILE.read_text(encoding="utf-8"))


@app.get("/api/subdirs")
async def api_subdirs():
    return JSONResponse({"subdirs": list_subdirs()})


@app.post("/api/mkdir")
async def api_mkdir(subdir: str = Form(...)):
    target = resolve_subdir(subdir)
    target.mkdir(parents=True, exist_ok=True)
    rel = str(target.relative_to(UPLOADS_ROOT)).replace("\\", "/")
    return JSONResponse({"ok": True, "path": rel})


# ── chunked upload ────────────────────────────────────────────────────────────

@app.post("/api/upload/start")
async def upload_start(
    filename: str = Form(...),
    subdir:   str = Form(""),
    total_size: int = Form(...),
):
    """Create a session, return session_id and how many bytes we already have."""
    fname  = safe_filename(filename)
    target = resolve_subdir(subdir)
    target.mkdir(parents=True, exist_ok=True)

    sid = str(uuid.uuid4())
    tmp = SESSIONS_ROOT / sid

    sessions[sid] = {
        "filename":   fname,
        "subdir":     subdir,
        "total_size": total_size,
        "received":   0,
        "tmp":        tmp,
    }

    return JSONResponse({"session_id": sid, "received": 0})


@app.post("/api/upload/chunk")
async def upload_chunk(
    session_id: str  = Form(...),
    chunk_index: int = Form(...),
    file: UploadFile = File(...),
):
    """Append one chunk to the temp file."""
    sid = safe_session_id(session_id)
    if sid not in sessions:
        raise HTTPException(404, "Session not found — may have expired")

    sess = sessions[sid]
    tmp: Path = sess["tmp"]

    data = await file.read()

    async with aiofiles.open(tmp, "ab") as f:
        await f.write(data)

    sess["received"] += len(data)

    return JSONResponse({
        "ok":       True,
        "received": sess["received"],
        "total":    sess["total_size"],
    })


@app.post("/api/upload/finish")
async def upload_finish(session_id: str = Form(...)):
    """Move assembled temp file to final destination."""
    sid = safe_session_id(session_id)
    if sid not in sessions:
        raise HTTPException(404, "Session not found")

    sess   = sessions[sid]
    tmp    = sess["tmp"]
    target = resolve_subdir(sess["subdir"])
    dest   = unique_dest(target / sess["filename"])

    target.mkdir(parents=True, exist_ok=True)

    # atomic move within same filesystem; fallback to copy+delete
    try:
        tmp.rename(dest)
    except OSError:
        import shutil
        shutil.move(str(tmp), str(dest))

    del sessions[sid]

    rel = str(dest.relative_to(UPLOADS_ROOT)).replace("\\", "/")
    return JSONResponse({"ok": True, "path": rel, "size": dest.stat().st_size})


@app.get("/api/upload/status/{session_id}")
async def upload_status(session_id: str):
    """Return how many bytes received so far (for resume on reconnect)."""
    sid = safe_session_id(session_id)
    if sid not in sessions:
        raise HTTPException(404, "Session not found")
    sess = sessions[sid]
    tmp: Path = sess["tmp"]
    received = tmp.stat().st_size if tmp.exists() else 0
    sessions[sid]["received"] = received
    return JSONResponse({"received": received, "total": sess["total_size"]})


@app.delete("/api/upload/{session_id}")
async def upload_cancel(session_id: str):
    """Abort and clean up a session."""
    sid = safe_session_id(session_id)
    if sid in sessions:
        tmp = sessions[sid]["tmp"]
        if tmp.exists():
            tmp.unlink()
        del sessions[sid]
    return JSONResponse({"ok": True})


# ── main ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import multiprocessing
    workers = min(multiprocessing.cpu_count(), 4)
    print(f"Upload server | uploads: {UPLOADS_ROOT} | http://localhost:8001")
    uvicorn.run(
        app, host="127.0.0.1", port=8001, workers=workers,
        timeout_keep_alive=3600, loop="asyncio",
    )