from __future__ import annotations

import dataclasses
import datetime
import enum
import functools
import json
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, TextIO

import structlog

from dissect.target.tools.dump.utils import (
    Compression,
    Serialization,
    get_current_utc_time,
    parse_datetime_iso,
)

if TYPE_CHECKING:
    from collections.abc import Iterator

    from typing_extensions import Self

    from dissect.target.target import Target
    from dissect.target.tools.dump.run import RecordStreamElement

log = structlog.get_logger(__name__)


STATE_FILE_NAME = "target-dump.state.json"

PENDING_UPDATES_LIMIT = 10


@dataclass
class Sink:
    target_path: str
    func: str
    path: Path
    is_dirty: bool = True
    record_count: int = 0
    size_bytes: int = 0

    def __post_init__(self):
        self.func = getattr(self.func, "name", self.func)


@dataclass
class DumpState:
    target_paths: list[str]
    functions: str
    excluded_functions: list[str]
    serialization: str
    compression: str
    start_time: datetime.datetime
    last_update_time: datetime.datetime

    sinks: list[Sink] = dataclasses.field(default_factory=list)

    # Volatile properties
    output_dir: Path | None = None
    pending_updates_count: int | None = 0

    @property
    def record_count(self) -> int:
        return sum(s.record_count for s in self.sinks)

    @property
    def finished_sinks(self) -> list[Sink]:
        return [sink for sink in self.sinks if not sink.is_dirty]

    @property
    def path(self) -> Path:
        return DumpState.get_state_path(self.output_dir)

    @classmethod
    def get_state_path(cls, output_dir: Path) -> Path:
        return output_dir / STATE_FILE_NAME

    def get_full_sink_path(self, sink: Sink) -> Path:
        if not self.output_dir:
            raise ValueError("Output directory is unknown for the state")
        return self.output_dir / sink.path

    def get_sink(self, path: Path) -> Sink | None:
        for sink in self.sinks:
            if sink.path == path:
                return sink
        return None

    def serialize(self) -> str:
        """Serialize state instance into a JSON formatted string."""
        state_dict = dataclasses.asdict(self)
        state_dict.pop("output_dir")
        state_dict.pop("pending_updates_count")
        return json.dumps(
            state_dict,
            default=serialize_obj,
            indent=4,
            sort_keys=True,
        )

    def persist(self, fh: TextIO) -> None:
        """Write serialized state instance into profided ``fh`` byte stream, overwriting it from the beginning."""
        fh.seek(0)
        fh.write(self.serialize())
        fh.flush()
        self.pending_updates_count = 0
        log.debug("State flushed")

    def mark_as_finished(self, target: Target, func: str) -> None:
        """Mark sinks that match provided ``target`` and ``func`` pair as not dirty."""
        matching_sinks = [
            sink for sink in self.sinks if str(sink.target_path) == str(target.path) and sink.func == func
        ]
        for sink in matching_sinks:
            sink.is_dirty = False

    def create_sink(self, sink_path: Path, stream_element: RecordStreamElement) -> Sink:
        """Create a sink instance for provided ``sink_path`` and ``stream_element``
        (from which ``target`` and ``func`` properties are used).
        """
        sink = Sink(
            path=sink_path,
            target_path=str(stream_element.target.path),
            func=stream_element.func,
        )
        self.sinks.append(sink)
        return sink

    def update(self, stream_element: RecordStreamElement, fp_position: int) -> None:
        """Update a sink instance for provided ``stream_element``."""
        sink = self.get_sink(stream_element.sink_path)

        sink.record_count += 1
        sink.size_bytes = fp_position
        self.last_update_time = get_current_utc_time()

        self.pending_updates_count += 1
        log.debug("State updated", records=self.record_count, sinks=len(self.sinks), path=self.path)

    @classmethod
    def from_dict(cls, state_dict: dict) -> Self:
        """Deserialize state instance from provided dictionary."""
        return cls(
            target_paths=state_dict["target_paths"],
            functions=state_dict["functions"],
            excluded_functions=state_dict["excluded_functions"],
            serialization=Serialization(state_dict["serialization"]),
            compression=Compression(state_dict["compression"]),
            start_time=parse_datetime_iso(state_dict["start_time"]),
            last_update_time=parse_datetime_iso(state_dict["last_update_time"]),
            sinks=[
                Sink(
                    target_path=sink["target_path"],
                    path=Path(sink["path"]),
                    func=sink["func"],
                    record_count=sink["record_count"],
                    size_bytes=sink["size_bytes"],
                    is_dirty=sink["is_dirty"],
                )
                for sink in state_dict["sinks"]
            ],
        )

    @classmethod
    def from_path(cls, output_dir: Path) -> Self | None:
        """Deserialize state instance from a file in the provided output directory path."""
        state_path = cls.get_state_path(output_dir)
        if not state_path.exists():
            return None

        with state_path.open(mode="r") as fh:
            try:
                state_dict = json.load(fh)
            except ValueError as e:
                log.warning("Can not load state from path", path=state_path, exc=e)
                return None

        state = cls.from_dict(state_dict)
        state.output_dir = output_dir
        return state

    def get_invalid_sinks(self) -> list[Sink]:
        """Return sinks that have a mismatch between recorded size and a real file size."""
        invalid_sinks = []
        for sink in self.sinks:
            # sink file does not exist
            if not self.get_full_sink_path(sink).exists():
                invalid_sinks.append(sink)
                continue

            # recorded file size for a clean sink is incorrect
            if not sink.is_dirty and sink.size_bytes != sink.path.stat().st_size:
                invalid_sinks.append(sink)
                continue

        return invalid_sinks

    def drop_invalid_sinks(self) -> None:
        """Remove sinks that have a mismatch between recorded size and a real file size from the list of sinks."""
        for invalid_sink in self.get_invalid_sinks():
            self.sinks.remove(invalid_sink)
            log.debug("Ignoring invalid sink", sink=invalid_sink.path)

    def drop_dirty_sinks(self) -> None:
        """Drop sinks that are marked as "dirty" in the current state from the list of sinks."""
        dirty_sinks = [s for s in self.sinks if s.is_dirty]
        for dirty_sink in dirty_sinks:
            self.sinks.remove(dirty_sink)
            log.debug("Ignoring dirty sink", sink=dirty_sink.path)


def create_state(
    *,
    output_dir: Path,
    target_paths: list[str],
    functions: str,
    excluded_functions: list[str],
    serialization: Serialization,
    compression: Compression = None,
) -> DumpState:
    """Create a ``DumpState`` instance with provided properties."""
    current_time = get_current_utc_time()
    return DumpState(
        target_paths=target_paths,
        functions=functions,
        excluded_functions=excluded_functions,
        serialization=serialization,
        compression=compression,
        start_time=current_time,
        last_update_time=current_time,
        output_dir=output_dir,
    )


@contextmanager
def persisted_state(state: DumpState) -> Iterator[Callable]:
    """Return a context manager for persisting ``DumpState`` instance."""

    def save_state(fh: TextIO) -> None:
        state.persist(fh)
        state.pending_updates_count = 0

    def save_state_when_over_limit(fh: TextIO) -> None:
        if state.pending_updates_count >= PENDING_UPDATES_LIMIT:
            save_state(fh)

    state.path.parent.mkdir(parents=True, exist_ok=True)

    with state.path.open(mode="w") as fh:
        try:
            yield functools.partial(save_state_when_over_limit, fh)
        finally:
            save_state(fh)


def load_state(output_dir: Path) -> DumpState | None:
    """Load persisted ``DumpState`` instance from provided ``output_dir`` path and perform sink validation."""
    state = DumpState.from_path(output_dir)

    if not state:
        return None

    # Dropping sinks that are marked as finished (clean) but the file
    # size on the disk is different from the one stored in the state object.
    state.drop_invalid_sinks()

    # Since there are no guarantees about the order of produced records
    # and the records themselves do not have unique IDs, we can not easily verify
    # what records are already stored in the dirty sink.
    # It is easier to just restart all dirty sinks.
    state.drop_dirty_sinks()

    return state


def serialize_obj(obj: Any) -> str:
    """JSON serializer for object types not serializable by ``json`` library."""
    if isinstance(obj, datetime.datetime):
        return obj.isoformat()
    if isinstance(obj, Path):
        return str(obj)
    if isinstance(obj, enum.Enum):
        return obj.value
    raise TypeError(f"Type {type(obj)} not serializable")
