import os
from datetime import datetime
from hashlib import sha256

import aiofiles
from Crypto.Random import get_random_bytes
from base58 import b58encode

from app.core._config import UPLOADS_DIR
from app.core._crypto.cipher import AESCipher
from app.core.models.keys import KnownKey
from app.core.models.node_storage import StoredContent
from app.core.logger import make_log
from base58 import b58decode


async def create_new_encryption_key(db_session, user_id: int = None) -> KnownKey:
    randpart = get_random_bytes(32)
    new_seed = randpart
    new_seed_str = b58encode(new_seed).decode()
    new_seed_hash_bin = sha256(new_seed).digest()
    new_seed_hash = b58encode(new_seed_hash_bin).decode()
    public_key = get_random_bytes(32)  # not used yet, algo is symmetric
    public_key_str = b58encode(public_key).decode()
    public_key_hash_bin = sha256(public_key).digest()
    public_key_hash = b58encode(public_key_hash_bin).decode()

    new_key = KnownKey(
        type="CONTENT_ENCRYPTION_KEY",
        seed=new_seed_str,
        seed_hash=new_seed_hash,
        public_key=public_key_str,
        public_key_hash=public_key_hash,

        algo="AES256",
        meta={"I_user_id": user_id} if user_id else None,
        created=datetime.now()
    )
    from sqlalchemy import select
    db_session.add(new_key)
    await db_session.commit()
    new_key = (await db_session.execute(select(KnownKey).where(KnownKey.seed_hash == new_seed_hash))).scalars().first()
    assert new_key, "Key not created"
    return new_key


async def create_encrypted_content(
    db_session, decrypted_content: StoredContent,
) -> StoredContent:
    from sqlalchemy import select
    # Try to find an already created encrypted counterpart for this decrypted content
    encrypted_content = (
        await db_session.execute(
            select(StoredContent).where(StoredContent.decrypted_content_id == decrypted_content.id)
        )
    ).scalars().first()
    if encrypted_content:
        make_log("create_encrypted_content", f"(d={decrypted_content.cid.serialize_v2()}) => (e={encrypted_content.cid.serialize_v2()}): already exist (found by decrypted content)", level="debug")
        return encrypted_content

    encrypted_content = None
    # Avoid accessing relationship attributes in async context to prevent MissingGreenlet
    if not decrypted_content.key_id:
        key = await create_new_encryption_key(db_session, user_id=decrypted_content.user_id)
        decrypted_content.key_id = key.id
        await db_session.commit()
        assert decrypted_content.key_id, "Key not assigned"

    # Explicitly load the key to avoid lazy-loading via relationship in async mode
    key = (
        await db_session.execute(select(KnownKey).where(KnownKey.id == decrypted_content.key_id))
    ).scalars().first()
    # If the referenced key is missing or malformed, create a fresh one
    if not key or not key.seed:
        key = await create_new_encryption_key(db_session, user_id=decrypted_content.user_id)
        decrypted_content.key_id = key.id
        await db_session.commit()

    decrypted_path = os.path.join(UPLOADS_DIR, decrypted_content.hash)
    decrypted_bin = b58decode(decrypted_content.hash)

    cipher = AESCipher(key.seed_bin)

    encrypted_bin = cipher.encrypt(decrypted_bin)
    encrypted_hash_bin = sha256(encrypted_bin).digest()
    encrypted_hash = b58encode(encrypted_hash_bin).decode()
    encrypted_content = (await db_session.execute(select(StoredContent).where(StoredContent.hash == encrypted_hash))).scalars().first()
    if encrypted_content:
        make_log("create_encrypted_content", f"(d={decrypted_content.cid.serialize_v2()}) => (e={encrypted_content.cid.serialize_v2()}): already exist (found by encrypted_hash)", level="debug")
        return encrypted_content

    encrypted_content = None

    encrypted_meta = dict(decrypted_content.meta or {})
    encrypted_meta["encrypt_algo"] = "AES256"

    encrypted_content = StoredContent(
        type="local/content_bin",
        hash=encrypted_hash,
        onchain_index=None,
        filename=decrypted_content.filename,
        meta=encrypted_meta,
        user_id=decrypted_content.user_id,

        encrypted=True,
        decrypted_content_id=decrypted_content.id,
        key_id=decrypted_content.key_id,

        created=datetime.now(),
    )
    db_session.add(encrypted_content)
    await db_session.commit()

    encrypted_path = os.path.join(UPLOADS_DIR, encrypted_hash)
    async with aiofiles.open(encrypted_path, mode='wb') as file:
        await file.write(encrypted_bin)

    encrypted_content = (await db_session.execute(select(StoredContent).where(StoredContent.hash == encrypted_hash))).scalars().first()
    assert encrypted_content, "Content not created"
    make_log("create_encrypted_content", f"(d={decrypted_content.cid.serialize_v2()}) => (e={encrypted_content.cid.serialize_v2()}): created new content/bin", level="debug")
    return encrypted_content