from __future__ import annotations

from collections.abc import Iterable
from datetime import datetime
from typing import Any, Protocol, NamedTuple
from pathlib import Path
from py_compile import PycInvalidationMode

import marshal
import struct


class PyCacheFile(NamedTuple):
    code: CodeObject
    magic: bytes  # 4 bytes
    invalidation_mode: PycInvalidationMode
    source_hash: bytes | None
    local_modification_date: datetime | None  # localtime
    file_size_bytes: bytes | None

    @property
    def magic_number(self) -> int:
        return struct.unpack("<H", self.magic[:2])[0]

    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, PyCacheFile):
            return False
        return self.code.co_code == other.code.co_code


# See https://github.com/python/typeshed/blob/main/stdlib/types.pyi#L122 for python-version adjusted
# type stub
class CodeObject(Protocol):
    # See https://docs.python.org/3/reference/datamodel.html#index-60
    co_name: str
    co_qualname: str
    co_argcount: int
    co_posonlyargcount: int
    co_kwonlyargcount: int
    co_nlocals: int
    co_varnames: tuple[str]
    co_cellvars: tuple[str]
    co_freevars: tuple[str]
    co_code: bytes
    co_consts: tuple
    co_names: tuple[str]
    co_filename: str
    co_firstlineno: int
    co_lnotab: str  # deprecated
    co_stacksize: int
    co_flags: int

    def co_positions(self) -> Iterable[Position]: ...

    def co_lins(self) -> Iterable[Line]: ...

    def replace(self, **kwargs) -> CodeObject: ...


class Position(NamedTuple):
    start_line: int
    end_line: int
    start_col: int
    end_col: int


class Line(NamedTuple):
    start: int
    end: int
    lineno: int


# inspired by https://github.com/nedbat/coveragepy/blob/master/lab/show_pyc.py
# TODO: this only works for "new" versions of python, perhaps for 3.7+
def parse_pyc_file(path: Path) -> PyCacheFile:
    # see PEP 3147
    mod_dt = None
    size = None
    source_hash = None
    with path.open("rb") as fp:
        magic_number = fp.read(4)
        flags = struct.unpack("<L", fp.read(4))[0]
        hash_based = bool(flags & 0x01)  # otherwise timestamp based
        check_source = bool(flags & 0x02)
        if hash_based:
            source_hash = fp.read(8)
        else:
            mod_date_bytes = fp.read(4)
            mod_dt = datetime.fromtimestamp(struct.unpack("<L", mod_date_bytes)[0])
            size_bytes = fp.read(4)
            size = struct.unpack("<L", size_bytes)[0]
        code = marshal.load(fp)  # this is at offset 16 bytes
    if hash_based:
        if check_source:
            invalidation_mode = PycInvalidationMode.CHECKED_HASH
        else:
            invalidation_mode = PycInvalidationMode.UNCHECKED_HASH
    else:
        invalidation_mode = PycInvalidationMode.TIMESTAMP
    return PyCacheFile(code, magic_number, invalidation_mode, source_hash, mod_dt, size)


if __name__ == "__main__":
    f = parse_pyc_file(Path("./__pycache__/models.cpython-313.pyc"))
    print(f)
    print(f == f)
    print(f.magic_number)