import hashlib
import json
import os
import shutil
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import parse_qs, urlparse


ROOT = Path(__file__).resolve().parent
PUBLIC_DIR = ROOT / "public"
UPLOAD_DIR = ROOT / "uploads"
CHUNK_DIR = UPLOAD_DIR / "chunks"
MERGED_DIR = UPLOAD_DIR / "merged"
PORT = int(os.environ.get("PORT", "8000"))


def ensure_dirs() -> None:
    CHUNK_DIR.mkdir(parents=True, exist_ok=True)
    MERGED_DIR.mkdir(parents=True, exist_ok=True)


def safe_upload_id(upload_id: str) -> str:
    return "".join(char for char in upload_id if char.isalnum() or char in "-_")[:120]


def json_response(handler: SimpleHTTPRequestHandler, status: int, payload: dict) -> None:
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json; charset=utf-8")
    handler.send_header("Content-Length", str(len(body)))
    handler.end_headers()
    handler.wfile.write(body)


def file_sha256(file_path: Path) -> str:
    digest = hashlib.sha256()
    with file_path.open("rb") as file:
        for block in iter(lambda: file.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


class ChunkUploadHandler(SimpleHTTPRequestHandler):
    def translate_path(self, path: str) -> str:
        request_path = urlparse(path).path
        if request_path == "/":
            request_path = "/index.html"
        return str(PUBLIC_DIR / request_path.lstrip("/"))

    def do_GET(self) -> None:
        parsed = urlparse(self.path)
        if parsed.path == "/api/status":
            self.handle_status(parsed.query)
            return

        super().do_GET()

    def do_POST(self) -> None:
        parsed = urlparse(self.path)
        if parsed.path == "/api/upload-chunk":
            self.handle_upload_chunk()
            return

        if parsed.path == "/api/merge":
            self.handle_merge()
            return

        json_response(self, 404, {"message": "not found"})

    def handle_status(self, query: str) -> None:
        params = parse_qs(query)
        upload_id = safe_upload_id(params.get("uploadId", [""])[0])
        if not upload_id:
            json_response(self, 400, {"message": "uploadId is required"})
            return

        upload_dir = CHUNK_DIR / upload_id
        uploaded = []
        if upload_dir.exists():
            uploaded = sorted(int(path.name) for path in upload_dir.iterdir() if path.name.isdigit())

        json_response(self, 200, {"uploadId": upload_id, "uploadedChunks": uploaded})

    def handle_upload_chunk(self) -> None:
        upload_id = safe_upload_id(self.headers.get("X-Upload-Id", ""))
        chunk_index = self.headers.get("X-Chunk-Index", "")
        total_chunks = self.headers.get("X-Total-Chunks", "")
        content_length = int(self.headers.get("Content-Length", "0"))

        if not upload_id or not chunk_index.isdigit() or not total_chunks.isdigit():
            json_response(self, 400, {"message": "invalid upload headers"})
            return

        upload_dir = CHUNK_DIR / upload_id
        upload_dir.mkdir(parents=True, exist_ok=True)
        chunk_path = upload_dir / chunk_index

        with chunk_path.open("wb") as file:
            remaining = content_length
            while remaining > 0:
                data = self.rfile.read(min(1024 * 1024, remaining))
                if not data:
                    break
                file.write(data)
                remaining -= len(data)

        json_response(self, 200, {
            "message": "chunk uploaded",
            "uploadId": upload_id,
            "chunkIndex": int(chunk_index),
            "totalChunks": int(total_chunks),
        })

    def handle_merge(self) -> None:
        content_length = int(self.headers.get("Content-Length", "0"))
        body = self.rfile.read(content_length)
        payload = json.loads(body.decode("utf-8") or "{}")
        upload_id = safe_upload_id(payload.get("uploadId", ""))
        file_name = Path(payload.get("fileName", "uploaded.bin")).name
        total_chunks = int(payload.get("totalChunks", 0))
        expected_hash = payload.get("fileHash", "")

        if not upload_id or total_chunks <= 0:
            json_response(self, 400, {"message": "invalid merge payload"})
            return

        upload_dir = CHUNK_DIR / upload_id
        missing = [index for index in range(total_chunks) if not (upload_dir / str(index)).exists()]
        if missing:
            json_response(self, 409, {"message": "chunks missing", "missingChunks": missing})
            return

        target_path = MERGED_DIR / f"{upload_id}-{file_name}"
        with target_path.open("wb") as target:
            for index in range(total_chunks):
                with (upload_dir / str(index)).open("rb") as chunk:
                    shutil.copyfileobj(chunk, target)

        actual_hash = file_sha256(target_path)
        if expected_hash and actual_hash != expected_hash:
            target_path.unlink(missing_ok=True)
            json_response(self, 422, {"message": "file hash mismatch", "actualHash": actual_hash})
            return

        shutil.rmtree(upload_dir, ignore_errors=True)
        json_response(self, 200, {
            "message": "file merged",
            "fileName": target_path.name,
            "size": target_path.stat().st_size,
            "sha256": actual_hash,
        })


if __name__ == "__main__":
    ensure_dirs()
    server = ThreadingHTTPServer(("127.0.0.1", PORT), ChunkUploadHandler)
    print(f"Chunk upload demo running at http://127.0.0.1:{PORT}")
    server.serve_forever()