# -*- coding: utf-8 -*-
import os
from typing import Optional, Union

from .base import SimpleSlice
from .data_slice import DataSlice
from v123pan.crypto.chacha20_encryptor import ChaCha20Encryptor


class ChaCha20Poly1305Slice(SimpleSlice):
    """
    生产级 XChaCha20-Poly1305 分片加密实现
    
    数据布局：
    - Header (44B): Salt(16) + Nonce(24) + BlockSize(4)
    - Body: [Encrypted Block (N bytes) + TAG (16 bytes)] * Block Count
    
    与 CloudDrive 完全兼容。
    """

    HEADER_SIZE = 44  # 16 + 24 + 4
    TAG_SIZE = 16
    DEFAULT_BLOCK_SIZE = 1024 * 1024  # 1MB

    def __init__(self, inner: SimpleSlice, password: str, block_size: int = DEFAULT_BLOCK_SIZE) -> None:
        super().__init__()
        if block_size <= 0:
            raise ValueError(f"Invalid block_size: {block_size}")

        self._inner = inner
        self._block_size = block_size
        self._password = password

        # 延迟计算加密后的总长度（在第一次读取时计算）
        self._length: Optional[int] = None
        self._block_count: Optional[int] = None

        self._position = 0
        self._salt: Optional[bytes] = None
        self._nonce: Optional[bytes] = None
        self._encryptor: Optional[ChaCha20Encryptor] = None

    @classmethod
    async def create(
        cls,
        data: Union[bytes, SimpleSlice],
        password: str,
        block_size: int = DEFAULT_BLOCK_SIZE
    ) -> 'ChaCha20Poly1305Slice':
        """
        创建加密分片

        Args:
            data: 字节数据或 SimpleSlice 对象
            password: 加密密码
            block_size: 分块大小

        Returns:
            ChaCha20Poly1305Slice: 加密分片实例

        Example:
            >>> async with await ChaCha20Poly1305Slice.create(b"data", "password") as slice:
            ...     encrypted = await slice.read_all()
        """
        if isinstance(data, bytes):
            inner = DataSlice(data)
        elif isinstance(data, SimpleSlice):
            inner = data
        else:
            raise TypeError(f"不支持的数据类型：{type(data)}")

        return cls(inner, password, block_size)
    
    def _calculate_length(self) -> None:
        """计算加密后的总长度"""
        if self._length is not None:
            return
        
        inner_len = len(self._inner)
        self._block_count = (inner_len + self._block_size - 1) // self._block_size
        self._length = self.HEADER_SIZE + inner_len + (self._block_count * self.TAG_SIZE)
    
    def __len__(self) -> int:
        """获取加密后的总长度"""
        self._calculate_length()
        return self._length

    def _init_encryptor(self) -> None:
        """初始化加密器并派生密钥"""
        if self._encryptor:
            return
        
        self._salt = os.urandom(16)
        # XChaCha20 需要 24 字节 nonce
        self._nonce = os.urandom(24) 
        
        self._encryptor = ChaCha20Encryptor(self._password, self._salt)
        self._encryptor.derive_key()

    def get_header(self) -> bytes:
        self._init_encryptor()
        if self._salt is None or self._nonce is None:
            raise ValueError("Encryptor not initialized")
        
        header = bytearray(self.HEADER_SIZE)
        header[0:16] = self._salt
        header[16:40] = self._nonce  # 24 字节 nonce
        header[40:44] = self._block_size.to_bytes(4, 'little')
        return bytes(header)

    def _virtual_to_raw(self, virtual_pos: int) -> int:
        """将加密流的位置映射回原始数据位置"""
        if virtual_pos <= self.HEADER_SIZE:
            return 0
        
        relative_pos = virtual_pos - self.HEADER_SIZE
        layout_unit = self._block_size + self.TAG_SIZE
        
        block_num = relative_pos // layout_unit
        block_offset = relative_pos % layout_unit
        
        # 如果 offset 落在 TAG 区域，映射到该块末尾
        actual_offset = min(block_offset, self._block_size)
        return block_num * self._block_size + actual_offset

    def _raw_to_virtual(self, raw_pos: int) -> int:
        """将原始位置映射到加密流位置"""
        block_num = raw_pos // self._block_size
        block_offset = raw_pos % self._block_size
        return self.HEADER_SIZE + block_num * (self._block_size + self.TAG_SIZE) + block_offset

    async def read_abstract(self, size: int) -> bytes:
        if self._encryptor is None:
            self._init_encryptor()

        # 边界检查
        if self._position >= self._length or size <= 0:
            return b""
        
        actual_size = min(size, self._length - self._position)
        result = bytearray()
        
        # 1. 处理 Header 读取
        if self._position < self.HEADER_SIZE:
            header = self.get_header()
            chunk_size = min(actual_size, self.HEADER_SIZE - self._position)
            result.extend(header[self._position : self._position + chunk_size])
            self._position += chunk_size
            actual_size -= chunk_size

        # 2. 处理主体加密数据读取
        if actual_size > 0:
            body_data = await self._read_encrypted_stream(self._position, actual_size)
            result.extend(body_data)
            self._position += len(body_data)

        return bytes(result)

    async def _read_encrypted_stream(self, v_start: int, size: int) -> bytes:
        """
        处理极其复杂的流式读取：跨块、非对齐读取。
        """
        layout_unit = self._block_size + self.TAG_SIZE
        v_relative = v_start - self.HEADER_SIZE
        
        # 计算涉及到的第一个块和最后一个块
        first_block_idx = v_relative // layout_unit
        last_block_idx = (v_relative + size - 1) // layout_unit
        
        # 确定需要读取的原始数据范围（按块对齐读取以保证解密正确）
        raw_read_start = first_block_idx * self._block_size
        raw_read_end = min((last_block_idx + 1) * self._block_size, len(self._inner))
        
        # 执行底层读取
        await self._inner.seek(raw_read_start)
        all_raw_data = await self._inner.read(raw_read_end - raw_read_start)
        
        # 对读取的块进行全量加密和加签
        temp_buffer = bytearray()
        for i in range(0, len(all_raw_data), self._block_size):
            curr_block_idx = first_block_idx + (i // self._block_size)
            chunk = all_raw_data[i : i + self._block_size]

            # 生成该块专属 Nonce (前 8 字节为 block_num)
            block_nonce = bytearray(self._nonce)
            block_nonce[0:8] = curr_block_idx.to_bytes(8, 'little')

            # 加密 (ChaCha20Poly1305 已经包含 16 字节 tag)
            encrypted_chunk = self._encryptor.encrypt(chunk, bytes(block_nonce))
            temp_buffer.extend(encrypted_chunk)

        # 从生成的全量块数据中，截取出用户请求的那一部分
        # 计算 v_start 在 temp_buffer 中的相对起始位置
        offset_in_buffer = v_relative % layout_unit
        return bytes(temp_buffer[offset_in_buffer : offset_in_buffer + size])

    async def _close(self) -> None:
        if self._inner:
            await self._inner.close()
            self._inner = None

    async def read_all(self) -> bytes:
        """
        读取全部加密数据

        Returns:
            bytes: 完整的加密数据（包含 Header）
        """
        await self.seek(0)
        return await self.read(None)