import os
import subprocess
from pathlib import Path
from typing import Optional
import psutil
from ._protocols import MountInfo, DiskInfo

class SystemAPI:

    def query_mounts(self) -> list[MountInfo]:
        mounts = []
        for partition in psutil.disk_partitions(all=False):
            mount_info = MountInfo(device=partition.device, mount_point=partition.mountpoint, fs_type=partition.fstype, options=partition.opts)
            mounts.append(mount_info)
        return mounts

    def query_nvmes(self) -> list[DiskInfo]:
        disks = []
        try:
            result = subprocess.run(['lsblk', '-ndo', 'NAME,MODEL,SIZE,SERIAL', '-b'], capture_output=True, text=True, check=False)
            if result.returncode == 0:
                for line in result.stdout.strip().split('\n'):
                    if not line.strip():
                        continue
                    parts = line.split(maxsplit=3)
                    if len(parts) >= 3:
                        device = f'/dev/{parts[0]}'
                        model = parts[1] if len(parts) > 1 else 'Unknown'
                        size_str = parts[2] if len(parts) > 2 else '0'
                        serial = parts[3] if len(parts) > 3 else 'Unknown'
                        try:
                            size = int(size_str)
                        except ValueError:
                            size = 0
                        disk_info = DiskInfo(device=device, model=model, size=size, serial=serial)
                        disks.append(disk_info)
        except FileNotFoundError:
            pass
        if not disks:
            disks = self._query_disks_fallback()
        return disks

    def _query_disks_fallback(self) -> list[DiskInfo]:
        disks = []
        seen_devices = set()
        for partition in psutil.disk_partitions(all=True):
            device = partition.device
            if not device.startswith('/dev/'):
                continue
            base_device = self._get_base_device(device)
            if base_device in seen_devices:
                continue
            seen_devices.add(base_device)
            try:
                usage = psutil.disk_usage(partition.mountpoint)
                size = usage.total
            except (PermissionError, OSError):
                size = 0
            disk_info = DiskInfo(device=base_device, model='Unknown', size=size, serial='Unknown')
            disks.append(disk_info)
        return disks

    def _get_base_device(self, device: str) -> str:
        if 'nvme' in device:
            if 'p' in device:
                return device.rsplit('p', 1)[0]
            return device
        import re
        match = re.match('(/dev/[a-z]+)', device)
        if match:
            return match.group(1)
        return device

    def get_disk_for_path(self, path: Path) -> Optional[str]:
        path = path.resolve()
        best_match = None
        best_match_len = 0
        for partition in psutil.disk_partitions():
            mount_point = Path(partition.mountpoint)
            try:
                if path == mount_point or mount_point in path.parents:
                    mount_len = len(str(mount_point))
                    if mount_len > best_match_len:
                        best_match = partition.device
                        best_match_len = mount_len
            except (ValueError, OSError):
                continue
        return best_match

    def get_disk_usage(self, path: Path) -> tuple[int, int, int]:
        try:
            usage = psutil.disk_usage(str(path))
            return (usage.total, usage.used, usage.free)
        except (PermissionError, OSError):
            return (0, 0, 0)

    def get_mount_point(self, path: Path) -> Optional[Path]:
        path = path.resolve()
        best_match = None
        best_match_len = 0
        for partition in psutil.disk_partitions():
            mount_point = Path(partition.mountpoint)
            try:
                if path == mount_point or mount_point in path.parents:
                    mount_len = len(str(mount_point))
                    if mount_len > best_match_len:
                        best_match = mount_point
                        best_match_len = mount_len
            except (ValueError, OSError):
                continue
        return best_match

    def is_same_filesystem(self, path1: Path, path2: Path) -> bool:
        try:
            stat1 = path1.stat()
            stat2 = path2.stat()
            return stat1.st_dev == stat2.st_dev
        except (OSError, PermissionError):
            return False