from __future__ import annotations

import asyncio

from typing import TYPE_CHECKING, Iterable, override

if TYPE_CHECKING:
    from jw.pkg.lib.file_types import FileType

from jw.pkg.lib.base import StatResult
from jw.pkg.lib.FcWrapper import FcWrapper
from jw.pkg.lib.file_types import FT_DIR, FT_REG
from jw.pkg.lib.FileContext import FileContext
from jw.pkg.lib.Result import Result

class FakeCtx(FileContext):
    """An in-memory FileContext that counts its primitive calls"""

    def __init__(self, uri: str, chroot: bool = False) -> None:
        super().__init__(uri, chroot = chroot)
        self.files: dict[str, bytes] = {}
        self.dirs: set[str] = set()
        self.counts: dict[str, int] = {}
        self.opens = 0

    def count(self, op: str) -> None:
        self.counts[op] = self.counts.get(op, 0) + 1

    def _stat_of(self, path: str) -> StatResult:
        if path in self.dirs:
            return StatResult(0o40755, '0', '0', 40, 0.0, 0.0, 0.0)
        if path in self.files:
            return StatResult(0o100644, '0', '0', len(self.files[path]), 0.0, 0.0, 0.0)
        raise FileNotFoundError(path)

    @override
    async def _open(self) -> None:
        self.opens += 1

    @override
    async def _close(self) -> None:
        pass

    @override
    async def _get(
        self, path: str, wd: str | None, throw: bool, verbose: bool | None, title: str
    ) -> Result:
        self.count('get')
        if wd is not None:
            path = wd + '/' + path
        if path not in self.files:
            raise FileNotFoundError(path)
        return Result(stdout = self.files[path])

    @override
    async def _put(
        self,
        path: str,
        content: bytes,
        wd: str | None,
        throw: bool,
        verbose: bool | None,
        title: str,
        owner: str | None,
        group: str | None,
        mode: str | None,
        atomic: bool,
        ftype: FileType,
    ) -> Result:
        self.count('put')
        if wd is not None:
            path = wd + '/' + path
        self.files[path] = content
        return Result(stdout = b'')

    @override
    async def _mkdir(self, path: str, mode: int, ftype: FileType) -> None:
        self.count('mkdir')
        self.dirs.add(path)

    @override
    async def _unlink(self, path: str) -> None:
        self.count('unlink')
        self.files.pop(path, None)
        self.dirs.discard(path)

    @override
    async def _stat(self, path: str, follow_symlinks: bool) -> StatResult:
        self.count('stat')
        return self._stat_of(path)

    @override
    async def _file_exists(self, path: str) -> bool:
        self.count('file_exists')
        try:
            self._stat_of(path)
            return True
        except FileNotFoundError:
            return False

    @override
    async def _is_dir(self, path: str, follow_symlinks: bool) -> bool:
        self.count('is_dir')
        try:
            st = self._stat_of(path)
            return (st.mode & 0o170000) == 0o40000
        except FileNotFoundError:
            return False

    @override
    async def _readdir(self, path: str) -> Iterable[str]:
        self.count('readdir')
        prefix = path.rstrip('/') + '/'
        return [
            p[len(prefix):] for p in sorted(self.files.keys() | self.dirs)
            if p.startswith(prefix)
        ]

async def main() -> None:
    fake = FakeCtx('local:///tmp/fc-wrapper-fake', chroot = True)
    fc = FcWrapper(fake)

    # -- The wrapper adopts the identity of the wrapped context
    assert fc.wrapped is fake
    assert fc.uri is fake.uri
    assert fc.root == fake.root
    assert fc.chroot == fake.chroot
    assert fc.username == fake.username
    assert fc.interactive == fake.interactive
    assert fc.verbose_default == fake.verbose_default

    # -- Put and get are forwarded; paths are composed once
    await fc.put('etc/a.conf', b'x=1\n', ftype = FT_REG)
    assert fake.files == {'/tmp/fc-wrapper-fake/etc/a.conf': b'x=1\n'}
    assert fake.counts['put'] == 1

    ret = await fc.get('etc/a.conf')
    assert ret.stdout == b'x=1\n'
    assert fake.counts['get'] == 1

    # -- stat, file_exists, is_dir, and readdir are forwarded
    st = await fc.stat('etc/a.conf')
    assert st.size == 4
    assert await fc.file_exists('etc/a.conf') is True
    assert await fc.file_exists('etc/missing.conf') is False
    await fc.mkdir('etc/sub', 0o755, FT_DIR)
    assert await fc.is_dir('etc/sub') is True
    names = await fc.readdir('etc')
    assert 'a.conf' in names and 'sub' in names

    # -- Unlink is forwarded
    await fc.unlink('etc/a.conf')
    assert '/tmp/fc-wrapper-fake/etc/a.conf' not in fake.files
    assert fake.counts['unlink'] == 1

    # -- Open and close reach the wrapped context's public methods
    await fc.open()
    await fc.close()
    assert fake.opens == 1

asyncio.run(main())