from PDAPI import *
import os
import json
import hashlib
import shutil
import time

RESOURCE_ROOT = os.path.join(os.path.abspath(os.path.join(os.getcwd(), "resource")))
TEMP_ROOT = os.path.join(RESOURCE_ROOT, ".tmp")


def _sha1_path(sha1_hex):
    return os.path.join(RESOURCE_ROOT, sha1_hex[0:2], sha1_hex[2:4], sha1_hex[4:6], sha1_hex)


def _ensure_dir(path):
    os.makedirs(os.path.dirname(path), exist_ok=True)


def _format_size(size):
    if size < 1024:
        return f"{size} B"
    elif size < 1024 * 1024:
        return f"{size / 1024:.1f} KB"
    elif size < 1024 * 1024 * 1024:
        return f"{size / (1024 * 1024):.1f} MB"
    else:
        return f"{size / (1024 * 1024 * 1024):.1f} GB"


def _cleanup_temp(temp_dir):
    try:
        shutil.rmtree(temp_dir, ignore_errors=True)
    except:
        pass


__metadata__ = {
    "methods": ["POST"],
    "blueprint": "GERP",
    "permission": 3.0,
    "attributes": [],
    "limit": "10 per minute",
}

def HandleRequest(_CleanedData):
    current_user = UGM.GetELTUser()
    if not current_user:
        return {"code": 401, "msg": "未登录"}, 401

    data = flask.request.get_json(silent=True) or {}
    upload_id = data.get("uploadId", "").strip()

    if not upload_id:
        return {"code": 1, "msg": "缺少uploadId"}, 400

    temp_dir = os.path.join(TEMP_ROOT, upload_id)
    if not os.path.isdir(temp_dir):
        return {"code": 1, "msg": "上传会话不存在或已过期"}, 404

    meta_path = os.path.join(temp_dir, "meta.json")
    if not os.path.exists(meta_path):
        _cleanup_temp(temp_dir)
        return {"code": 1, "msg": "上传元数据不存在"}, 404

    with open(meta_path, "r", encoding="utf-8") as f:
        meta = json.load(f)

    total_chunks = meta["totalChunks"]
    received = sorted(meta.get("receivedChunks", []))
    expected = list(range(total_chunks))

    if received != expected:
        missing = [i for i in expected if i not in received]
        return {"code": 1, "msg": f"分片不完整，缺少: {missing}"}, 400

    sha1 = hashlib.sha1()
    merged_size = 0
    final_path = os.path.join(temp_dir, "merged")

    with open(final_path, "wb") as out_f:
        for i in range(total_chunks):
            chunk_path = os.path.join(temp_dir, str(i))
            if not os.path.exists(chunk_path):
                _cleanup_temp(temp_dir)
                return {"code": 1, "msg": f"分片{i}文件丢失"}, 400
            with open(chunk_path, "rb") as in_f:
                while True:
                    block = in_f.read(65536)
                    if not block:
                        break
                    sha1.update(block)
                    out_f.write(block)
                    merged_size += len(block)

    sha1_hex = sha1.hexdigest()

    existing = GERPFileAsset.query.filter(GERPFileAsset.FileSHA1 == sha1_hex).first()
    if existing:
        _cleanup_temp(temp_dir)
        return {
            "code": 0,
            "msg": "文件已存在(秒传)",
            "data": {
                "assetUUID": UUIDObj(existing.AssetUUID).UpperStr,
                "fileName": existing.FileName,
                "fileSize": existing.FileSize,
                "fileSizeDisplay": _format_size(existing.FileSize),
                "mimeType": existing.MimeType,
                "sha1": existing.FileSHA1,
                "refLink": f"https://api.gvsds.com/Asset/Download/{existing.FileSHA1}",
                "createTime": existing.CreateTime.strftime("%Y-%m-%d %H:%M:%S") if existing.CreateTime else None,
            }
        }, 200

    storage_path = _sha1_path(sha1_hex)
    _ensure_dir(storage_path)

    try:
        shutil.move(final_path, storage_path)
    except shutil.Error:
        with open(final_path, "rb") as src, open(storage_path, "wb") as dst:
            shutil.copyfileobj(src, dst)
        os.remove(final_path)

    for i in range(total_chunks):
        chunk_path = os.path.join(temp_dir, str(i))
        if os.path.exists(chunk_path):
            os.remove(chunk_path)

    asset_uuid = NewBUUID()
    new_asset = GERPFileAsset()
    new_asset.AssetUUID = asset_uuid
    new_asset.FileName = meta["fileName"]
    new_asset.FileSHA1 = sha1_hex
    new_asset.FileSize = merged_size
    new_asset.MimeType = meta["mimeType"]
    new_asset.StoragePath = storage_path
    new_asset.UploaderUUID = current_user.UserUUID
    new_asset.RefCount = 0

    current_role_uuid_str = flask.request.cookies.get("CurrentRoleUUID")
    if current_role_uuid_str:
        try:
            role_uuid_bytes = UUIDObj(current_role_uuid_str).bytes
            current_role = GERPRole.query.filter(GERPRole.RoleUUID == role_uuid_bytes).first()
            if current_role and current_role.CompanyUUID:
                new_asset.CompanyUUID = current_role.CompanyUUID
        except:
            pass

    try:
        db.session.add(new_asset)
        db.session.commit()
    except:
        db.session.rollback()
        if os.path.exists(storage_path):
            os.remove(storage_path)
        _cleanup_temp(temp_dir)
        return {"code": 1, "msg": "数据库写入失败"}, 200

    _cleanup_temp(temp_dir)

    return {
        "code": 0,
        "msg": "上传成功",
        "data": {
            "assetUUID": UUIDObj(asset_uuid).UpperStr,
            "fileName": new_asset.FileName,
            "fileSize": new_asset.FileSize,
            "fileSizeDisplay": _format_size(new_asset.FileSize),
            "mimeType": new_asset.MimeType,
            "sha1": new_asset.FileSHA1,
            "refLink": f"https://api.gvsds.com/Asset/Download/{new_asset.FileSHA1}",
            "createTime": new_asset.CreateTime.strftime("%Y-%m-%d %H:%M:%S") if new_asset.CreateTime else None,
        }
    }, 201