import logging
from typing import Iterator, List, Optional, Tuple

from volatility3.framework import exceptions, interfaces, renderers
from volatility3.framework.configuration import requirements
from volatility3.framework.renderers import format_hints
from volatility3.plugins.windows import bigpools

vollog = logging.getLogger(__name__)


class BigPoolDump(interfaces.plugins.PluginInterface):


    _version = (0, 1, 0)
    _required_framework_version = (2, 0, 0)

    @classmethod
    def get_requirements(cls) -> List[interfaces.configuration.RequirementInterface]:
        return [
            requirements.ModuleRequirement(
                name="kernel",
                description="Windows kernel",
                architectures=["Intel32", "Intel64"],
            ),
            requirements.VersionRequirement(
                name="bigpools", component=bigpools.BigPools, version=(2, 0, 0)
            ),
            requirements.StringRequirement(
                name="tags",
                description="Comma-separated list of pool tags to dump (e.g. TcDN,Tcpt)",
                optional=False,
            ),
            requirements.BooleanRequirement(
                name="include-free",
                description="Include freed allocations",
                default=False,
                optional=True,
            ),
        ]

    def _dump_one(self, addr: int, size: int, tag: str) -> Optional[str]:
        kernel = self.context.modules[self.config["kernel"]]
        layer = self.context.layers[kernel.layer_name]

        filename = f"bigpool.{tag}.0x{addr:016x}.0x{size:x}.dmp"
        try:
            data = layer.read(addr, size, pad=True)
        except exceptions.InvalidAddressException:
            return None

        try:
            with self.open(filename) as fp:
                fp.write(data)
            return filename
        except OSError:
            return None

    def _generator(self) -> Iterator[Tuple[int, Tuple[object, ...]]]:
        tags = [t.strip() for t in (self.config.get("tags") or "").split(",") if t.strip()]
        if not tags:
            vollog.warning("No tags specified")
            return

        show_free = bool(self.config.get("include-free"))

        for big_pool in bigpools.BigPools.list_big_pools(
            context=self.context,
            kernel_module_name=self.config["kernel"],
            tags=tags,
            show_free=show_free,
        ):
            tag = big_pool.get_key()
            size = big_pool.get_number_of_bytes()
            if isinstance(size, interfaces.renderers.BaseAbsentValue):
                continue
            addr = int(big_pool.Va) & ~1 

            dumped_as = self._dump_one(addr, int(size), tag)
            if dumped_as is None:
                dumped_as = renderers.UnreadableValue()

            status = "Free" if big_pool.is_free() else "Allocated"

            yield (
                0,
                (
                    format_hints.Hex(addr),
                    tag,
                    format_hints.Hex(int(size)),
                    big_pool.get_pool_type(),
                    status,
                    dumped_as,
                ),
            )

    def run(self) -> renderers.TreeGrid:
        return renderers.TreeGrid(
            [
                ("Allocation", format_hints.Hex),
                ("Tag", str),
                ("NumberOfBytes", format_hints.Hex),
                ("PoolType", str),
                ("Status", str),
                ("File output", str),
            ],
            self._generator(),
        )