"""
hdlg - Modern cross-platform GUI for hdl-dump.
Copyright (C) 2021 rlaphoenix

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""

from __future__ import annotations

import struct
import subprocess
import sys
from ctypes.wintypes import HANDLE
from pathlib import Path
from typing import Union

import win32con
import win32file
import winioctlcon

from hdlg.utils import NEIGHBORING_WHITESPACE


class HDD:
    def __init__(self, target: Union[str, Path], model: str):
        self.handle = win32file.INVALID_HANDLE_VALUE
        self.target = target
        self.model = model

        self._geometry = None
        self._disk_size = None
        self._is_apa_partitioned = None
        self._apa_checksum = None

        self.open(target)

    def __enter__(self):
        return self

    def __exit__(self, *_, **__):
        self.dispose()

    def dispose(self):
        if self.handle != win32file.INVALID_HANDLE_VALUE:
            win32file.CloseHandle(self.handle)

    def open(self, device: str, extended=False) -> HANDLE:
        if sys.platform == "win32":
            if not device.startswith("\\\\.\\"):
                device = r"\\.\%s" % device  # unc target
        self.handle = win32file.CreateFile(
            # https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilea
            device,  # target
            win32con.MAXIMUM_ALLOWED,  # 0x0080 | 0x0020,  # desired access
            win32con.FILE_SHARE_READ | win32con.FILE_SHARE_WRITE,  # share mode
            None,  # security attributes
            win32con.OPEN_EXISTING,  # creation disposition
            win32con.FILE_ATTRIBUTE_NORMAL,  # flags and attributes
            None  # template file
        )
        if self.handle == win32file.INVALID_HANDLE_VALUE:
            raise RuntimeError("Failed to obtain device handle...")
        if extended:
            # TODO: Sometimes works, sometimes doesn't. Only got it working (once) on a CD-ROM drive with a DVD.
            print(win32file.DeviceIoControl(self.handle, winioctlcon.FSCTL_ALLOW_EXTENDED_DASD_IO, None, None))
        return self.handle

    def seek(self, to: int, whence: int = win32file.FILE_BEGIN) -> int:
        pos = win32file.SetFilePointer(self.handle, to, whence)
        if pos != to:
            raise IOError(f"Seek was not precise...")
        return pos

    def read(self, size: int) -> bytes:
        if size % 512 != 0:
            raise ValueError("Size must be a multiple of 512 for some reason, Ask Windows.")
        res, data = win32file.ReadFile(self.handle, size, None)
        if res != 0:
            raise IOError(f"An error occurred: {res} {data}")
        if len(data) < size:
            raise IOError(f"Read {size - len(data)} less bytes than requested...")
        return data

    @property
    def geometry(self) -> tuple[int, ...]:
        """
        Retrieves information about the physical disk's geometry.
        https://docs.microsoft.com/en-us/windows/win32/api/winioctl/ns-winioctl-disk_geometry_ex

        Returns a tuple of:
            Cylinders-Lo
            Cylinders-Hi
            Media Type
            Tracks Per Cylinder
            Sectors Per Track
            Bytes Per Sector
            Disk Size
            Extra Data
        """
        if self._geometry is not None:
            return self._geometry

        if self.handle == win32file.INVALID_HANDLE_VALUE:
            raise ValueError("Handle is invalid, have you opened a handle yet?")

        self._geometry = struct.unpack("8L", win32file.DeviceIoControl(
            self.handle,  # handle
            winioctlcon.IOCTL_DISK_GET_DRIVE_GEOMETRY_EX,  # ioctl api
            None,  # in buffer
            32  # out buffer
        ))
        return self._geometry

    @property
    def disk_size(self) -> int:
        """Get full Disk Size (in bytes)."""
        if self._disk_size is not None:
            return self._disk_size

        cyl_lo, cyl_hi, _, tpc, spt, bps, _, _ = self.geometry

        self._disk_size = (cyl_lo + cyl_hi) * tpc * spt * bps
        return self._disk_size

    @property
    def is_apa_partitioned(self) -> bool:
        """Check if HDD is a PS2 APA-formatted device."""
        if self._is_apa_partitioned is not None:
            return self._is_apa_partitioned

        old_pos = self.seek(0, whence=win32file.FILE_CURRENT)

        try:
            if self.seek(0) != 0:
                raise RuntimeError("Unable to seek to start of HDD, cannot check if APA Partitioned...")
            header = self.read(1024)
            checksum = header[0:4]
            magic = header[4:8]

            new_checksum = struct.pack(
                "<Q",
                sum(struct.unpack("<I", header[n:n+4])[0] for n in range(4, 1024, 4))
            )[:4]

            self._is_apa_partitioned = magic == b"APA\0" and checksum == new_checksum
        finally:
            self.seek(old_pos)

        return self._is_apa_partitioned

    @property
    def apa_checksum(self):
        if not self.is_apa_partitioned:
            return None

        if self._apa_checksum is not None:
            return self._apa_checksum

        old_pos = self.seek(0, whence=win32file.FILE_CURRENT)

        try:
            if self.seek(0) != 0:
                raise RuntimeError("Unable to seek to start of HDD, cannot get APA checksum...")
            header = self.read(512)
            self._apa_checksum = header[0:4]
        finally:
            self.seek(old_pos)

        return self._apa_checksum

    def get_games_list(self) -> list[tuple[str, int, int, str, str, str]]:
        """
        Get a list of games installed on the HDD (if any).

        Returns a tuple of:
            MediaType
            Size (in Bytes)
            Flags
            DMA
            GameID
            GameName
        """
        # TODO: Add handler for when there's no games installed
        hdl_toc = subprocess.check_output([
            "hdl_dump.exe",
            "hdl_toc",
            self.target.upper().replace(r"\\.\PHYSICALDRIVE", "hdd") + ":"
        ])
        games = hdl_toc.decode().splitlines(keepends=False)[1:-1]
        games = [
            NEIGHBORING_WHITESPACE.sub(" ", game).split(" ", maxsplit=5)
            for game in games
        ]
        games = [
            (media_type, int(size.replace("KB", "")) * 1000, int(flags), dma, game_id, game_name)
            for media_type, size, flags, dma, game_id, game_name in games
        ]
        return games
