from __future__ import annotations

import gzip
import logging
import os
import re
import stat
import time
import zlib
from pathlib import Path
from typing import Optional, List, TypeVar, Type, Tuple

from concurrent_log_handler import ConcurrentRotatingFileHandler
from tarka.utility.utf8 import is_utf8_codepoint_start

logger = logging.getLogger(__name__)

_LogFileRefT = TypeVar("_LogFileRefT")


class _LogFileRef:
    __slots__ = ("path", "stat", "index")

    @classmethod
    def create(cls: Type[_LogFileRefT], reader: ConcurrentRotatingFileReader, path: Path) -> Optional[_LogFileRefT]:
        m = reader.name_re.match(path.name)
        if not m:
            return
        index = int(m.group(1)) if m.group(1) else 0
        try:
            s = path.lstat()
        except OSError:
            return
        if stat.S_ISREG(s.st_mode):
            return cls(path, s, index)

    def __init__(self, p: Path, s: os.stat_result, i: int):
        self.path = p
        self.stat = s
        self.index = i

    @property
    def stat_identity(self) -> Tuple[int, int]:
        return self.stat.st_dev, self.stat.st_ino

    def stat_matches(self, other: _LogFileRef) -> bool:
        return self.stat.st_dev == other.stat.st_dev and self.stat.st_ino == other.stat.st_ino

    @property
    def valid(self) -> bool:
        try:
            s = self.path.lstat()
        except OSError:
            return False
        return self.stat.st_dev == s.st_dev and self.stat.st_ino == s.st_ino

    def read_bytes(self) -> bytes:
        data = self.path.read_bytes()
        if self.path.name.endswith(".gz"):
            return gzip.decompress(data)
        return data

    def read_bytes_range(self, from_: int, limit: int) -> bytes:
        with open(self.path, "rb") as f:
            f.seek(from_)
            return f.read(limit)


class _ReaderState:
    __slots__ = (
        "refs",
        "forward_index",
        "forward_pos",
        "forward_cache",
        "backward_index",
        "backward_pos",
        "backward_cache",
    )

    def __init__(self):
        self.refs: List[Optional[_LogFileRef]] = []
        self.forward_index: int = 0
        self.forward_pos: int = 0
        self.forward_cache: Tuple[int, int, bytes] = (0, 0, b"")
        self.backward_index: int = 0
        self.backward_pos: int = 0
        self.backward_cache: Tuple[int, int, bytes] = (0, 0, b"")

    def is_new(self) -> bool:
        return not self.refs

    def refs_need_update(self) -> bool:
        return (
            self.is_new()
            or (
                0 <= self.forward_index < len(self.refs)
                and (not self.refs[self.forward_index] or not self.refs[self.forward_index].valid)
            )
            or (
                0 <= self.backward_index < len(self.refs)
                and (not self.refs[self.backward_index] or not self.refs[self.backward_index].valid)
            )
        )

    def copy(self) -> _ReaderState:
        s = _ReaderState()
        s.refs = list(self.refs)
        s.forward_index = self.forward_index
        s.forward_pos = self.forward_pos
        s.forward_cache = self.forward_cache
        s.backward_index = self.backward_index
        s.backward_pos = self.backward_pos
        s.backward_cache = self.backward_cache
        return s


class ConcurrentRotatingFileReader(ConcurrentRotatingFileHandler):
    WINDOW_MIN_LENGTH = 16384
    NEWLINE_ALLOWANCE = 4096

    @classmethod
    def create_file_name_re(cls, handler: ConcurrentRotatingFileHandler) -> re.Pattern:
        prefix = re.escape(Path(handler.baseFilename).name)
        suffix = re.escape(".gz" if handler.use_gzip else "")
        return re.compile(rf"^{prefix}(?:\.(\d+){suffix})?$")

    @classmethod
    def slice_logs_forward(cls, logs: bytes, from_: int, step: int) -> Tuple[bytes, int]:
        assert step >= 8
        from_ = max(0, min(len(logs), from_))
        to = min(len(logs), from_ + step)  # limit to length
        if to == len(logs):
            return logs[from_:], to
        nl = logs.rfind(b"\n", max(from_, to - cls.NEWLINE_ALLOWANCE), to)
        if from_ <= nl:
            return logs[from_ : nl + 1], nl + 1
        for i in range(to, max(from_ + 1, to - 5), -1):
            if is_utf8_codepoint_start(logs[i]):
                return logs[from_:i], i
        return logs[from_:to], to

    @classmethod
    def slice_logs_backward(cls, logs: bytes, from_: int, step: int) -> Tuple[bytes, int]:
        assert step >= 8
        from_ = max(0, min(len(logs), from_))
        to = max(0, from_ - step)  # limit to length
        if to == 0:
            return logs[:from_], to
        nl = logs.find(b"\n", to - 1, min(to + cls.NEWLINE_ALLOWANCE, from_ - 1))
        if 0 <= nl:
            return logs[nl + 1 : from_], nl + 1
        for i in range(to, min(from_ - 1, to + 5)):
            if is_utf8_codepoint_start(logs[i]):
                return logs[i:from_], i
        return logs[to:from_], to

    def __init__(self, handler: ConcurrentRotatingFileHandler, max_size: int):
        ConcurrentRotatingFileHandler.__init__(
            self, handler.baseFilename, encoding=handler.encoding, use_gzip=handler.use_gzip
        )  # NOTE: Other attrs are not applicable.
        self.lockFilename = handler.lockFilename  # necessary in case of custom lock_file_directory
        self.maxLockAttempts *= 3

        self.name_re = self.create_file_name_re(handler)
        self._max_size = max_size
        self._window_size = max(max_size // 2, self.WINDOW_MIN_LENGTH)

        self._state = _ReaderState()

    def read(self, forward: bool = True) -> dict:
        for i in range(5):
            try:
                s = self._state.copy()

                # Update the file references, this will usually just work, but race condition with log-rotation is
                # possible, so we always doublecheck for consistency.
                new = s.is_new()
                if s.refs_need_update():
                    self._manage_refs(s)
                result = self._read_forward(s) if forward else self._read_backward(s)
                check_refs = self._read_refs()
                if len(s.refs) != len(check_refs) or any(
                    (a is None) is not (b is None) or (a and not a.stat_matches(b)) for a, b in zip(s.refs, check_refs)
                ):
                    raise OSError("Log files check inconsistency")

                if new:
                    result["new"] = True
                self._state = s
                return result
            except OSError:
                if i == 4:
                    raise
                time.sleep(0.013)

    def _manage_refs(self, s: _ReaderState):
        new_refs = self._read_refs()
        if s.refs:
            self._track_refs(s, new_refs)
        else:
            self._begin_refs(s, new_refs)
        s.refs = new_refs

    def _read_refs(self) -> List[Optional[_LogFileRef]]:
        # collect a sorted list of log files from disk
        new_refs = []
        for item in Path(self.baseFilename).parent.iterdir():
            ref = _LogFileRef.create(self, item)
            if ref:
                new_refs.append(ref)
        new_refs.sort(key=lambda r: r.index)
        # ensure starting point and sequential validity, add None head if it is missing
        if new_refs:
            if new_refs[0].index == 0:
                pass
            elif new_refs[0].index == 1:
                new_refs.insert(0, None)
            else:
                new_refs[:] = [None]
            for i in range(1, len(new_refs)):
                if new_refs[i].index != i:
                    del new_refs[i:]
                    break
        else:
            new_refs.append(None)
        return new_refs

    def _begin_refs(self, s: _ReaderState, new_refs: List[_LogFileRef]):
        """
        Initialise the f/b index and pos fields.
        """
        if new_refs[0]:
            index = 0  # start in the current/head log
        elif len(new_refs) <= 1:
            index = 0  # no logs now at all, wait on first index
        else:
            index = 1  # start in the first archive
        pos = 0
        if 0 <= index < len(new_refs) and new_refs[index]:
            logs = new_refs[index].read_bytes()
            pos = max(0, len(logs) - self._window_size)
            while 0 < pos < len(logs) and not is_utf8_codepoint_start(logs[pos]):
                pos += 1
        s.forward_index = s.backward_index = index
        s.forward_pos = s.backward_pos = pos

    def _track_refs(self, s: _ReaderState, new_refs: List[_LogFileRef]):
        """
        Adjust the f/b index fields to the possibly rotated log files.
        """
        # NOTE: the rotating implementation of the superclass ensures that the head-ref will be invalidated first
        assert s.refs
        if len(s.refs) > 1:
            # try and get index delta from the move of the first archive (it's identity is stable unlike the head's)
            old_1_identity = s.refs[1].stat_identity
            for i in range(1, len(new_refs)):
                if new_refs[i].stat_identity == old_1_identity:
                    index_delta = i - 1
                    break
            else:
                index_delta = 3 * max(len(s.refs), len(new_refs))  # over-rotated
        else:
            # when there was only or no head before, we'll track it to the rightmost place
            index_delta = len(new_refs) - 1
        s.forward_index += index_delta
        s.backward_index += index_delta

    def _read_forward(self, s: _ReaderState):
        if s.forward_index >= len(s.refs):
            return {"logs": b"", "reset": True}  # can't continue
        ref = s.refs[s.forward_index]
        if s.forward_index == 0:  # read head log
            s.forward_cache = (0, 0, b"")
            if ref:
                logs, read_bytes_len = self.slice_logs_forward(
                    ref.read_bytes_range(s.forward_pos, self._window_size), 0, self._window_size
                )
                s.forward_pos += read_bytes_len  # slicing is done from 0 after seek-read
            else:  # waiting on missing head
                logs = b""
        else:  # read from an archive
            si = ref.stat_identity
            if si != s.forward_cache[:2]:
                s.forward_cache = (si[0], si[1], ref.read_bytes())
            logs, s.forward_pos = self.slice_logs_forward(s.forward_cache[2], s.forward_pos, self._window_size)
            if len(s.forward_cache[2]) == s.forward_pos:
                s.forward_index -= 1
                s.forward_pos = 0
        return {"logs": logs}

    def _read_backward(self, s: _ReaderState):
        c = zlib.compressobj(wbits=zlib.MAX_WBITS)
        size = 0
        chunks = []
        result = {"compression": "deflate"}
        while True:
            self._cache_backward(s)
            if s.backward_index >= len(s.refs) or not s.refs[s.backward_index]:
                result["end"] = True
                chunks.append(c.flush())
                break
            logs, s.backward_pos = self.slice_logs_backward(s.backward_cache[2], s.backward_pos, self._window_size)
            chunk = c.compress(bytes(reversed(logs)))
            size += len(chunk)
            chunks.append(chunk)
            if size >= self._window_size:
                chunks.append(c.flush())
                break
            chunk = c.flush(zlib.Z_SYNC_FLUSH)
            size += len(chunk)
            chunks.append(chunk)
            if size >= self._window_size:
                chunks.append(c.flush())
                break
        result["logs"] = b"".join(chunks)
        return result

    def _cache_backward(self, s: _ReaderState):
        if s.backward_pos == 0:
            s.backward_index += 1
        if s.backward_index >= len(s.refs) or not s.refs[s.backward_index]:
            s.backward_cache = (0, 0, b"")
        else:
            ref = s.refs[s.backward_index]
            si = ref.stat_identity
            if si != s.backward_cache[:2]:
                s.backward_cache = (si[0], si[1], ref.read_bytes())
                s.backward_pos = len(s.backward_cache[2])
