""" Python bindings for the Windows OLE Storage API """

from __future__ import annotations

from collections.abc import Mapping
from ctypes import *
from ctypes.wintypes import *
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import cached_property
from io import RawIOBase
from typing import Generator, Literal, TYPE_CHECKING
if TYPE_CHECKING:
    from ctypes import _CData, _FuncPointer
    from typing_extensions import Self
    from _typeshed import StrPath

    ParamMode = Literal['in', 'out', 'inout']
    Param = tuple[ParamMode, type[_CData], str] | tuple[ParamMode, type[_CData], str, object]

ole32 = oledll.ole32

def bind_winfunc(ret: type[_CData] | None, loc: CDLL | int, symname: str, *params: Param) -> _FuncPointer:
    def _param(p: Param) -> tuple[int, str] | tuple[int, str, object]:
        if   p[0] == 'in': m = 1
        elif p[0] == 'out': m = 2
        elif p[0] == 'inout': m = 3
        else: raise ValueError()

        return (m, p[2]) if len(p) == 3 else (m, p[2], p[3])
    ftype = WINFUNCTYPE(ret, *(p[1] for p in params))
    pflags = tuple(_param(p) for p in params)
    if isinstance(loc, int):
        return ftype(loc, symname, pflags)
    else:
        return ftype((symname, loc), pflags)


class GUID(Structure):
    _fields_ = [
        ('', c_ulong),
        ('', c_ushort),
        ('', c_ushort),
        ('', c_ubyte * 8),
    ]

class STGOPTIONS(Structure):
    _fields_ = [
        ('version', USHORT),
        ('reserved', USHORT),
        ('sectorSize', ULONG),
        ('templateFile', c_wchar_p),
    ]

class STATSTG(Structure):
    _fields_ = [
        ('name', LPOLESTR),
        ('type', DWORD),
        ('size', ULARGE_INTEGER),
        ('mtime', FILETIME),
        ('ctime', FILETIME),
        ('atime', FILETIME),
        ('mode', DWORD),
        ('locksSupported', DWORD),
        ('clsid', GUID),
        ('state', DWORD),
        ('reserved', DWORD),
    ]

class STATPROPSTG(Structure):
    _fields_ = [
        ('name', LPOLESTR),
        ('propid', DWORD),
        ('vt', USHORT),
    ]

class PROPSPEC(Structure):
    class _UNION(Union):
        _fields_ = [
            ('propid', ULONG),
            ('name', LPOLESTR),
        ]

    _anonymous_ = ("union",)
    _fields_ = [
        ('kind', ULONG),
        ('union', _UNION),
    ]

class PROPVARIANT(Structure):
    _fields_ = [
        ('vt', USHORT),
        ('', BYTE * 2),
        ('', ULONG),
        ('data', c_ubyte * 8),
    ]

CLSIDFromString = bind_winfunc(
    HRESULT, ole32, 'CLSIDFromString',
    ('in', LPCWSTR, 'str'),
    ('out', POINTER(GUID), 'clsid'),
)

StringFromCLSID = bind_winfunc(
    HRESULT, ole32, 'StringFromCLSID',
    ('in', POINTER(GUID), 'clsid'),
    ('out', POINTER(OLESTR), 'str'),
)

IIDFromString = bind_winfunc(
    HRESULT, ole32, 'IIDFromString',
    ('in', LPCWSTR, 'str'),
    ('out', POINTER(GUID), 'iid'),
)

CoInitializeEx = bind_winfunc(
    HRESULT, ole32, 'CoInitializeEx',
    ('in', c_void_p, 'reserved', None),
    ('in', DWORD, 'coinit', 0x00),
)

CoUninitialize = bind_winfunc(
    None, ole32, 'CoUninitialize'
)

CoCreateInstance = bind_winfunc(
    HRESULT, ole32, 'CoCreateInstance',
    ('in', POINTER(GUID), 'clsid'),
    ('in', c_void_p, 'outer', None),
    ('in', DWORD, 'clsContext'),
    ('in', POINTER(GUID), 'iid'),
    ('out', POINTER(LPVOID), 'cls'),
)

CoTaskMemFree = bind_winfunc(
    None, ole32, 'CoTaskMemFree',
    ('in', LPVOID, 'ptr'),
)

StgOpenStorageEx = bind_winfunc(
    HRESULT, ole32, 'StgOpenStorageEx',
    ('in', c_wchar_p, 'name'),
    ('in', DWORD, 'mode'),
    ('in', DWORD, 'fmt'),
    ('in', DWORD, 'attrs'),
    ('inout', POINTER(STGOPTIONS), 'options'),
    ('in', c_void_p, 'securityDescriptor', None),
    ('in', POINTER(GUID), 'iid'),
    ('out', POINTER(c_void_p), 'objectOpen'),
)

StgOpenPropStg = bind_winfunc(
    HRESULT, ole32, 'StgOpenPropStg',
    ('in', LPVOID, 'stg'),
    ('in', POINTER(GUID), 'fmtid'),
    ('in', DWORD, 'flags', 0x00),
    ('in', DWORD, 'reserved1', 0x00),
    ('out', POINTER(LPVOID), 'propstg'),
)

IUnknown_Release = bind_winfunc(
    c_ulong, 2, 'Release'
)

IStorage_EnumElements = bind_winfunc(
    HRESULT, 11, 'EnumElements',
    ('in', DWORD, 'reserved1', 0x00),
    ('in', c_void_p, 'reserved2', None),
    ('in', DWORD, 'reserved3', 0x00),
    ('out', POINTER(c_void_p), 'enum'),
)

IStorage_Stat = bind_winfunc(
    HRESULT, 17, 'Stat',
    ('out', POINTER(STATSTG), 'statstg'),
    ('in', DWORD, 'statFlag', 0x00),
)

IStorage_OpenStorage = bind_winfunc(
    HRESULT, 6, 'OpenStorage',
    ('in', c_wchar_p, 'name'),
    ('in', c_void_p, 'priority', None),
    ('in', DWORD, 'mode'),
    ('in', c_void_p, 'exclude', None),
    ('in', DWORD, 'reserved', 0x00),
    ('out', POINTER(LPVOID), 'stg'),
)

IStorage_OpenStream = bind_winfunc(
    HRESULT, 4, 'OpenStream',
    ('in', c_wchar_p, 'name'),
    ('in', c_void_p, 'reserved1', None),
    ('in', DWORD, 'mode'),
    ('in', DWORD, 'reserved2', 0x00),
    ('out', POINTER(LPVOID), 'stm'),
)

IEnumSTATSTG_Next = bind_winfunc(
    HRESULT, 3, 'Next',
    ('in', ULONG, 'count', 1),
    ('out', POINTER(STATSTG), 'elt'),
    ('out', POINTER(ULONG), 'n'),
)

IStream_Stat = bind_winfunc(
    HRESULT, 12, 'Stat',
    ('out', POINTER(STATSTG), 'statstg'),
    ('in', DWORD, 'statFlag', 0x00),
)

IStream_Seek = bind_winfunc(
    HRESULT, 5, 'Seek',
    ('in', LARGE_INTEGER, 'offset'),
    ('in', DWORD, 'origin'),
    ('out', POINTER(ULARGE_INTEGER), 'newPosition'),
)

IStream_Clone = bind_winfunc(
    HRESULT, 13, 'Clone',
    ('out', POINTER(LPVOID), 'stm'),
)

ISequentialStream_Read = bind_winfunc(
    HRESULT, 3, 'Read',
    ('in', c_char_p, 'buf'),
    ('in', ULONG, 'count'),
    ('out', PULONG, 'n'),
)

IPropertyStorage_Enum = bind_winfunc(
    HRESULT, 11, 'Enum',
    ('out', POINTER(LPVOID), 'enum')
)

IPropertyStorage_ReadMultiple = bind_winfunc(
    HRESULT, 3, 'ReadMultiple',
    ('in', ULONG, 'count'),
    ('in', POINTER(PROPSPEC), 'propspecs'),
    ('out', POINTER(PROPVARIANT), 'props'),
)

IEnumSTATPROPSTG_Next = bind_winfunc(
    HRESULT, 3, 'Next',
    ('in', ULONG, 'count', 1),
    ('out', POINTER(STATPROPSTG), 'elt'),
    ('out', POINTER(ULONG), 'n'),
)


class COM:
    def __enter__(self) -> Self:
        CoInitializeEx()
        return self

    def __exit__(self, *_) -> None:
        CoUninitialize()

class COMObject:
    _ptr: LPVOID
    
    def __init__(self, ptr: LPVOID | int) -> None:
        self._ptr = cast(ptr, LPVOID)

    def __del__(self):
        if self._ptr: IUnknown_Release(self._ptr)

class Storage(COMObject, Mapping):
    @cached_property
    def _stats(self) -> STATSTG:
        return IStorage_Stat(self._ptr)
    
    @cached_property
    def _child_info(self) -> dict[str, STATSTG]:
        def _gen() -> Generator[STATSTG, None, None]:
            enumerator = cast(IStorage_EnumElements(self._ptr), LPVOID)
            while True:
                child, n = IEnumSTATSTG_Next(enumerator)
                if not n: break
                yield child
            IUnknown_Release(enumerator)
        return { stats.name: stats for stats in _gen() }
    
    @property
    def name(self) -> str:
        return self._stats.name
    
    @property
    def size(self) -> int:
        return self._stats.size
    
    def __str__(self):
        return f'<Storage "{self.name}", size {self.size}>'
    
    def __iter__(self):
        return iter(self._child_info)
    
    def __len__(self):
        return len(self._child_info)
    
    def __getitem__(self, key: str) -> Storage | Stream | PropSet:
        child_stats = self._child_info[key]
        if child_stats.type == 1:
            child_ptr = IStorage_OpenStorage(self._ptr, name=child_stats.name, mode=0x0000_0020)
            return Storage(child_ptr)
        elif child_stats.type == 2:
            child_ptr = IStorage_OpenStream(self._ptr, name=child_stats.name, mode=0x0000_0010)
            return Stream(child_ptr)
        else:
            raise RuntimeError(f'Unsupported storage type: {child_stats.type}')

class Stream(COMObject, RawIOBase):
    @cached_property
    def _stats(self) -> STATSTG:
        return IStream_Stat(self._ptr)
    
    @property
    def name(self) -> str:
        return self._stats.name
    
    @property
    def size(self) -> int:
        return self._stats.size
    
    def readable(self) -> bool:
        return True
    
    def writable(self) -> bool:
        return False
    
    def seekable(self) -> bool:
        return True
    
    def read(self, count: int = -1, /) -> bytes:
        if count < 0: return self.readall()

        buf = create_string_buffer(count)
        read_ct = ISequentialStream_Read(self._ptr, buf, count)
        return bytes(buf)[:read_ct]
    
    def readall(self) -> bytes:
        return self.read(self.size)
    
    def seek(self, offset: int, whence: int = 0) -> int:
        return IStream_Seek(self._ptr, offset, whence)
    
    def clone(self) -> Self:
        ptr = IStream_Clone(self._ptr)
        return Stream(ptr)
    
    def open_propset(self, fmtid: str | GUID) -> PropSet:
        if isinstance(fmtid, str): fmtid = IIDFromString(fmtid)
        if not isinstance(fmtid, GUID): raise TypeError()
        ptr = StgOpenPropStg(self._ptr, byref(fmtid))
        return PropSet(ptr)
    
    def __str__(self):
        return f'<Stream "{self.name}", size {self.size}>'

class PropSet(COMObject, Mapping):
    @dataclass
    class _PropInfo:
        names: dict[str, int]
        info: dict[int, STATPROPSTG]

        def __iter__(self):
            return iter((self.info[k].name or k) for k in self.info)
        
        def __len__(self):
            return len(self.info)

    @cached_property
    def _prop_info(self) -> _PropInfo:
        def _gen() -> Generator[STATPROPSTG, None, None]:
            enumerator = cast(IPropertyStorage_Enum(self._ptr), LPVOID)
            while True:
                child, n = IEnumSTATPROPSTG_Next(enumerator)
                if not n: break
                yield child
            IUnknown_Release(enumerator)
        info = { stats.propid: stats for stats in _gen() }
        names = { stats.name: stats.propid for stats in info.values() if stats.name }
        return self._PropInfo(names, info)
    
    def __str__(self):
        return f'<PropSet>'
    
    def __iter__(self):
        return iter(self._prop_info)
    
    def __len__(self):
        return len(self._prop_info)
    
    def __getitem__(self, key: str | int) -> str | int | datetime | None:
        if isinstance(key, str):
            key = self._prop_info.names[key]
        if key not in self._prop_info.info: raise IndexError()

        propspec = PROPSPEC(0x01, propid=key)
        prop = IPropertyStorage_ReadMultiple(self._ptr, 1, byref(propspec))

        # only the few prop types needed for SummaryInformation are supported
        if prop.vt <= 0x01:  # EMPTY | VOID
            return None
        elif prop.vt == 0x1e:  # LPSTR
            pstr = cast(int.from_bytes(prop.data[:sizeof(LPSTR)], 'little'), LPSTR)
            assert pstr.value is not None
            return str(pstr.value, encoding='ansi')
        elif prop.vt == 0x40:  # FILETIME
            ft = int.from_bytes(prop.data[:sizeof(FILETIME)], 'little')
            if ft == 0: return None
            return datetime.fromtimestamp(ft // 10_000_000 - 11_644_473_600, timezone.utc)
        elif prop.vt == 0x03:  # I4
            return int.from_bytes(prop.data[:4], 'little')
        else:
            raise RuntimeError(f'Unsupported property type: {prop.vt:02x}')

def open_storage(path: StrPath) -> Storage:
    IID_IStorage = IIDFromString('{0000000b-0000-0000-C000-000000000046}')
    options = STGOPTIONS(version=1, reserved=0, sectorSize=512, templateFile=None)

    _, ptr = StgOpenStorageEx(
        name=str(path),
        mode=0x0000_0020,
        fmt=0x05,
        attrs=0x00,
        options=options,
        iid=IID_IStorage)
    return Storage(ptr)
