from __future__ import annotations

import asyncio
import os
import tempfile

from typing import TYPE_CHECKING, override

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

from jw.pkg.lib.base import StatResult
from jw.pkg.lib.ec.Local import Local
from jw.pkg.lib.ExecContext import ExecContext
from jw.pkg.lib.FcCache import FcCache
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] = {}

    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 _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)

class CountingLocal(Local):
    """A Local context that counts its executed commands"""

    def __init__(self) -> None:
        super().__init__('local')
        self.runs = 0

    @override
    async def _run(
        self,
        cmd: list[str],
        wd: str | None,
        verbose: bool,
        cmd_input: bytes | None,
        mod_env: dict[str, str] | None,
        interactive: bool,
        log_prefix: str,
    ) -> Result:
        self.runs += 1
        return await super()._run(
            cmd, wd, verbose, cmd_input, mod_env, interactive, log_prefix
        )

async def main() -> None:
    fake = FakeCtx('local:///tmp/fc-cache-fake', chroot = True)
    fc = FcCache(fake)
    await fc.mkdir('d', 0o755, FT_DIR)
    await fc.put('d/a', b'x', ftype = FT_REG)

    # -- Successful stats are cached
    st = await fc.stat('d/a')
    assert st.size == 1
    assert (await fc.stat('d/a')).size == 1
    assert fake.counts['stat'] == 1

    # -- Absences are cached, too
    assert await fc.file_exists('d/missing') is False
    assert await fc.file_exists('d/missing') is False
    assert fake.counts['stat'] == 2

    # -- Gets are cached per (path, wd, throw)
    ret = await fc.get('d/a')
    assert (await fc.get('d/a')).stdout == b'x'
    assert fake.counts['get'] == 1
    ret = await fc._get('a', '/tmp/fc-cache-fake/d', True, None, 'get')
    assert fake.counts['get'] == 2
    assert ret.stdout == b'x'

    # -- Writes drop the cache entries of their paths
    await fc.put('d/a', b'xy', ftype = FT_REG)
    assert (await fc.get('d/a')).stdout == b'xy'
    assert fake.counts['get'] == 3
    assert (await fc.stat('d/a')).size == 2
    assert fake.counts['stat'] == 3

    # -- Unlink drops them as well
    await fc.unlink('d/a')
    assert await fc.file_exists('d/a') is False
    assert fake.counts['stat'] == 4
    assert '/tmp/fc-cache-fake/d/a' not in fake.files

    # -- The wrapper itself is not an exec context
    assert not isinstance(fc, ExecContext)

    # -- warm() on a non-exec context: one stat per path, then cached
    fake2 = FakeCtx('local:///tmp/fc-cache-warm', chroot = True)
    fc2 = FcCache(fake2)
    await fc2.put('p/q', b'z', ftype = FT_REG)
    await fc2.warm(['p', 'p/q', 'p/r'])
    assert fake2.counts['stat'] == 3
    await fc2.warm(['p', 'p/q', 'p/r'])
    assert fake2.counts['stat'] == 3
    assert await fc2.file_exists('p/r') is False
    assert fake2.counts['stat'] == 3

    # -- warm() on an exec context: one host command for all paths
    with tempfile.TemporaryDirectory() as td:
        f = os.path.join(td, 'there')
        with open(f, 'w') as fh:
            fh.write('hello\n')
        cl = CountingLocal()
        fc3 = FcCache(cl)
        missing = os.path.join(td, 'not-there')
        await fc3.warm([f, missing])
        assert cl.runs == 1

        st = await fc3.stat(f)
        assert cl.runs == 1
        assert st.size == 6
        assert await fc3.file_exists(missing) is False
        assert cl.runs == 1
        try:
            await fc3.stat(missing)
            assert False, 'stat of a warmed absence must raise'
        except FileNotFoundError:
            pass
        assert cl.runs == 1

        # -- A put through the cache lands in the wrapped context and
        #    drops its cache entries, so the next stat is a fresh one.
        #    Local stats without a command, so file_exists adds no run
        g = os.path.join(td, 'new')
        await fc3.put(g, b'fresh\n', ftype = FT_REG)
        assert cl.runs == 2
        assert (await fc3.get(g)).stdout == b'fresh\n'
        assert cl.runs == 3
        assert await fc3.file_exists(g) is True
        assert cl.runs == 3
        assert await fc3.file_exists(g) is True
        assert cl.runs == 3

asyncio.run(main())