import os
import uuid
import hashlib
import json
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

# Максимальный размер одного чанка (чуть больше CHUNK_SIZE для допуска)
MAX_CHUNK_BYTES = CHUNK_SIZE + 1024


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

def resolve_subdir(subdir: str) -> Path:
    """Безопасное разрешение поддиректории внутри UPLOADS_ROOT."""
    if not subdir or subdir.strip() in ("", "/", "."):
        return UPLOADS_ROOT
    # Убираем любые управляющие символы и нулевые байты
    clean_raw = subdir.strip().replace("\x00", "")
    # as_posix + lstrip убирает ведущие слэши и нормализует разделители
    clean = Path(clean_raw).as_posix().lstrip("/")
    if not clean or clean == ".":
        return UPLOADS_ROOT
    target = (UPLOADS_ROOT / clean).resolve()
    # Проверяем что итоговый путь строго внутри UPLOADS_ROOT
    uploads_resolved = UPLOADS_ROOT.resolve()
    try:
        target.relative_to(uploads_resolved)
    except ValueError:
        raise HTTPException(400, "Invalid subdirectory")
    return target


def safe_filename(name: str) -> str:
    """Безопасное имя файла: только basename, без нулей и служебных символов."""
    name = Path(name).name          # отсекаем любой путь
    name = name.replace("\x00", "").strip()
    # Убираем ведущие точки и пробелы (скрытые файлы / Windows-артефакты)
    name = name.lstrip(". ")
    if not name:
        raise HTTPException(400, "Invalid filename")
    # Запрещаем зарезервированные имена Windows
    reserved = {"CON", "PRN", "AUX", "NUL",
                "COM1","COM2","COM3","COM4","COM5","COM6","COM7","COM8","COM9",
                "LPT1","LPT2","LPT3","LPT4","LPT5","LPT6","LPT7","LPT8","LPT9"}
    if name.upper().split(".")[0] in reserved:
        raise HTTPException(400, "Reserved filename")
    return name


def safe_session_id(sid: str) -> str:
    """UUID только — никаких path-трюков."""
    try:
        return str(uuid.UUID(sid))
    except (ValueError, AttributeError):
        raise HTTPException(400, "Invalid session id")


def safe_chunk_index(idx: int, total_chunks: int) -> int:
    """Проверяем что индекс чанка в допустимом диапазоне."""
    if idx < 0 or idx >= total_chunks:
        raise HTTPException(400, f"chunk_index out of range [0, {total_chunks - 1}]")
    return idx


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


def sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


# ── session state ─────────────────────────────────────────────────────────────
# {
#   session_id: {
#     "filename":     str,
#     "subdir":       str,
#     "total_size":   int,
#     "total_chunks": int,
#     "session_dir":  Path,          # .sessions/<uuid>/
#     "hashes":       {              # chunk_index -> { "client": str, "server": str, "ok": bool }
#                       0: {...}, 1: {...}, ...
#                     }
#   }
# }
sessions: dict[str, dict] = {}


def session_dir(sid: str) -> Path:
    return SESSIONS_ROOT / sid


def chunk_path(sid: str, idx: int) -> Path:
    return session_dir(sid) / f"chunk_{idx:08d}"


def hashes_path(sid: str) -> Path:
    return session_dir(sid) / "hashes.json"


async def save_hashes(sid: str) -> None:
    """Записываем текущее состояние хешей в JSON-файл сессии (для диагностики)."""
    sess = sessions[sid]
    data = {
        "session_id":   sid,
        "filename":     sess["filename"],
        "total_size":   sess["total_size"],
        "total_chunks": sess["total_chunks"],
        "chunks":       sess["hashes"],
    }
    async with aiofiles.open(hashes_path(sid), "w", encoding="utf-8") as f:
        await f.write(json.dumps(data, indent=2, ensure_ascii=False))


# ── 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(...),
):
    """
    Создать новую сессию загрузки.
    Возвращает session_id, количество чанков и уже принятые индексы
    (при перезапуске на уже существующей сессии — пусто).
    """
    if total_size < 0 or total_size > 500 * 1024 * 1024 * 1024:  # 500 ГБ
        raise HTTPException(400, "Invalid total_size")

    fname  = safe_filename(filename)
    target = resolve_subdir(subdir)
    target.mkdir(parents=True, exist_ok=True)

    sid           = str(uuid.uuid4())
    sdir          = session_dir(sid)
    sdir.mkdir(parents=True, exist_ok=True)

    total_chunks  = max(1, -(-total_size // CHUNK_SIZE))   # ceil div

    sessions[sid] = {
        "filename":     fname,
        "subdir":       subdir,
        "total_size":   total_size,
        "total_chunks": total_chunks,
        "session_dir":  sdir,
        "hashes":       {},
    }

    await save_hashes(sid)

    return JSONResponse({
        "session_id":   sid,
        "total_chunks": total_chunks,
        "received_chunks": [],          # список уже принятых индексов (пустой для новой сессии)
    })


@app.post("/api/upload/chunk")
async def upload_chunk(
    session_id:   str  = Form(...),
    chunk_index:  int  = Form(...),
    client_hash:  str  = Form(...),     # SHA-256 hex от клиента
    file: UploadFile   = File(...),
):
    """
    Принять один чанк, проверить его хеш.
    При несовпадении возвращает ok=False — клиент должен переотправить.
    При совпадении сохраняет чанк в отдельном файле и записывает хеш в JSON.
    """
    sid = safe_session_id(session_id)
    if sid not in sessions:
        raise HTTPException(404, "Session not found — may have expired")

    sess = sessions[sid]

    # Валидируем индекс чанка
    safe_chunk_index(chunk_index, sess["total_chunks"])

    # Валидируем формат client_hash (SHA-256 = 64 hex символа)
    client_hash = client_hash.strip().lower()
    if len(client_hash) != 64 or not all(c in "0123456789abcdef" for c in client_hash):
        raise HTTPException(400, "Invalid client_hash format (expected SHA-256 hex)")

    # Читаем данные с ограничением размера
    data = await file.read(MAX_CHUNK_BYTES + 1)
    if len(data) > MAX_CHUNK_BYTES:
        raise HTTPException(413, "Chunk too large")

    # Считаем хеш на сервере
    server_hash = sha256_bytes(data)

    hashes_entry = {
        "client": client_hash,
        "server": server_hash,
        "ok":     server_hash == client_hash,
        "size":   len(data),
    }
    sess["hashes"][str(chunk_index)] = hashes_entry
    await save_hashes(sid)

    if server_hash != client_hash:
        # Чанк не сохраняем — просим переотправить
        return JSONResponse({
            "ok":          False,
            "chunk_index": chunk_index,
            "server_hash": server_hash,
            "client_hash": client_hash,
            "error":       "Hash mismatch — please resend this chunk",
        }, status_code=200)   # 200 чтобы клиент прочитал тело

    # Хеш совпал — сохраняем чанк в отдельный файл
    cpath = chunk_path(sid, chunk_index)
    async with aiofiles.open(cpath, "wb") as f:
        await f.write(data)

    # Считаем принятые чанки
    received_chunks = [
        int(k) for k, v in sess["hashes"].items() if v["ok"]
    ]

    return JSONResponse({
        "ok":              True,
        "chunk_index":     chunk_index,
        "server_hash":     server_hash,
        "received_chunks": received_chunks,
        "total_chunks":    sess["total_chunks"],
    })


@app.post("/api/upload/finish")
async def upload_finish(session_id: str = Form(...)):
    """
    Собрать все чанки в итоговый файл (по порядку индексов).
    Перед сборкой проверяем что все чанки приняты с верным хешем.
    """
    sid = safe_session_id(session_id)
    if sid not in sessions:
        raise HTTPException(404, "Session not found")

    sess         = sessions[sid]
    total_chunks = sess["total_chunks"]
    hashes       = sess["hashes"]

    # Проверяем что все чанки есть и все ок
    missing = []
    bad     = []
    for i in range(total_chunks):
        key = str(i)
        if key not in hashes:
            missing.append(i)
        elif not hashes[key]["ok"]:
            bad.append(i)

    if missing or bad:
        return JSONResponse({
            "ok":      False,
            "missing": missing,
            "bad":     bad,
            "error":   "Not all chunks received or verified",
        }, status_code=409)

    # Собираем файл
    target = resolve_subdir(sess["subdir"])
    dest   = unique_dest(target / sess["filename"])
    target.mkdir(parents=True, exist_ok=True)

    async with aiofiles.open(dest, "wb") as out:
        for i in range(total_chunks):
            cpath = chunk_path(sid, i)
            if not cpath.exists():
                raise HTTPException(500, f"Chunk file missing for index {i}")
            async with aiofiles.open(cpath, "rb") as inp:
                while True:
                    buf = await inp.read(1024 * 1024)
                    if not buf:
                        break
                    await out.write(buf)

    # Финализируем JSON с хешами рядом с итоговым файлом
    final_hashes_path = dest.with_suffix(dest.suffix + ".chunks.json")
    hashes_data = {
        "session_id":   sid,
        "filename":     sess["filename"],
        "total_size":   sess["total_size"],
        "total_chunks": total_chunks,
        "chunks": {
            k: v for k, v in sorted(hashes.items(), key=lambda x: int(x[0]))
        },
    }
    async with aiofiles.open(final_hashes_path, "w", encoding="utf-8") as f:
        await f.write(json.dumps(hashes_data, indent=2, ensure_ascii=False))

    # Удаляем временную папку сессии
    import shutil
    sdir = sess["session_dir"]
    try:
        shutil.rmtree(str(sdir))
    except Exception:
        pass

    del sessions[sid]

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


@app.get("/api/upload/status/{session_id}")
async def upload_status(session_id: str):
    """Вернуть список принятых и верифицированных чанков (для resume)."""
    sid = safe_session_id(session_id)
    if sid not in sessions:
        raise HTTPException(404, "Session not found")
    sess = sessions[sid]
    received_chunks = [
        int(k) for k, v in sess["hashes"].items() if v["ok"]
    ]
    return JSONResponse({
        "total_chunks":    sess["total_chunks"],
        "received_chunks": received_chunks,
        "total_size":      sess["total_size"],
    })


@app.delete("/api/upload/{session_id}")
async def upload_cancel(session_id: str):
    """Отменить сессию и почистить временные файлы."""
    sid = safe_session_id(session_id)
    if sid in sessions:
        import shutil
        sdir = sessions[sid]["session_dir"]
        try:
            shutil.rmtree(str(sdir))
        except Exception:
            pass
        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",
    )