# 경로/파일: realcom-nas-ui/modules/disk.py
# 목적: 디스크/파티션 + 사용량(df) + 온도/SMART(smartctl) 조회
# 붙여넣기: 이 파일 전체 교체

from __future__ import annotations

import json
import os
import re
import shutil
import subprocess
from typing import Any, Dict, List, Tuple

SMARTCTL_CANDIDATES = ["/usr/sbin/smartctl", "/sbin/smartctl", "smartctl"]


def _is_exec_file(p: str) -> bool:
    return os.path.exists(p) and os.path.isfile(p) and os.access(p, os.X_OK)


def _which_smartctl() -> str:
    for p in SMARTCTL_CANDIDATES:
        if p == "smartctl":
            w = shutil.which("smartctl")
            if w:
                return w
        else:
            if _is_exec_file(p):
                return p
    # 마지막 폴백
    return "smartctl"


SMARTCTL = _which_smartctl()


def _run(cmd: List[str], timeout: int = 6) -> Tuple[int, str]:
    """stdout+stderr 합쳐서 반환"""
    try:
        p = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            timeout=timeout,
            check=False,
        )
        return p.returncode, (p.stdout or "")
    except subprocess.TimeoutExpired:
        return 124, "timeout"
    except Exception as e:
        return 1, f"error: {e}"


def _run_smartctl(args: List[str], timeout: int = 6) -> Tuple[int, str]:
    """
    smartctl은 권한 이슈 많음.
    1) smartctl 직접
    2) 권한 문제면 sudo -n smartctl 재시도 (비번 요구하면 즉시 실패)
    """
    rc, out = _run([SMARTCTL] + args, timeout=timeout)
    if rc == 0:
        return rc, out

    lower = (out or "").lower()
    if ("permission denied" in lower) or ("requires root" in lower) or ("operation not permitted" in lower):
        # -n : password prompt 금지 (멈춤 방지)
        rc2, out2 = _run(["sudo", "-n", SMARTCTL] + args, timeout=timeout)
        return rc2, out2

    return rc, out


def get_disk_temperature(dev_path: str):
    """
    smartctl 실행 후 온도값 추출
    Returns: int 또는 "N/A"
    """
    dev_path = (dev_path or "").strip()
    if not dev_path.startswith("/dev/"):
        return "N/A"

    # NVMe 우선
    if "nvme" in dev_path:
        rc, out = _run_smartctl(["-a", dev_path], timeout=6)
        if rc == 0 and out:
            m = re.search(r"Temperature:\s+(\d+)\s+Celsius", out, re.I)
            if m:
                v = int(m.group(1))
                if 0 < v < 110:
                    return v
            m2 = re.search(r"Temperature Sensor \d+:\s+(\d+)\s*C", out, re.I)
            if m2:
                v = int(m2.group(1))
                if 0 < v < 110:
                    return v

    # ATA/SATA
    rc, out = _run_smartctl(["-A", dev_path], timeout=6)
    if rc == 0 and out:
        # 1) Temperature_Celsius 라인 우선
        for key in ("Temperature_Celsius", "Airflow_Temperature_Cel"):
            m = re.search(rf"^\s*\d+\s+{key}\s+.*$", out, re.I | re.M)
            if m:
                line = m.group(0)
                nums = re.findall(r"\b(\d{1,3})\b", line)
                for n in reversed(nums):  # 뒤에서부터 현실값 찾기
                    v = int(n)
                    if 0 < v < 110:
                        return v

        # 2) attribute id 190/194 기반(일부 디스크는 이름이 다름)
        for m in re.finditer(r"^\s*(190|194)\s+.+$", out, re.M):
            line = m.group(0)
            nums = re.findall(r"\b(\d{1,3})\b", line)
            for n in reversed(nums):
                v = int(n)
                if 0 < v < 110:
                    return v

        # 3) 최후 폴백
        for line in out.splitlines():
            if re.search(r"Temperature", line, re.I):
                nums = re.findall(r"\b(\d{1,3})\b", line)
                for n in nums:
                    v = int(n)
                    if 0 < v < 110:
                        return v

    return "N/A"


def get_disk_smart_status(dev_path: str) -> str:
    """
    smartctl -H 명령으로 SMART 상태 조회
    Returns: "정상", "오류", "지원안함", "알수없음", "N/A"
    """
    dev_path = (dev_path or "").strip()
    if not dev_path.startswith("/dev/"):
        return "N/A"

    rc, out = _run_smartctl(["-H", dev_path], timeout=6)
    if rc != 0 and not out:
        return "N/A"

    if "PASSED" in out:
        return "정상"
    if "FAILED" in out:
        return "오류"

    lower = (out or "").lower()
    if "smart support is: unavailable" in lower or "smart disabled" in lower or "unknown usb bridge" in lower:
        return "지원안함"

    return "알수없음"


def _df_map() -> Dict[str, Dict[str, str]]:
    """df 결과를 /dev/xxx 키로 매핑"""
    rc, out = _run(["df", "-h", "--output=source,size,used,avail,pcent,target"], timeout=6)
    if rc != 0 or not out:
        return {}

    lines = out.strip().splitlines()
    if len(lines) <= 1:
        return {}

    mp: Dict[str, Dict[str, str]] = {}
    for line in lines[1:]:
        cols = line.split()
        if len(cols) != 6:
            continue
        src = cols[0].strip()
        if not src.startswith("/dev/"):
            continue
        key = src.replace("/dev/", "")
        mp[key] = {
            "size": cols[1],
            "used": cols[2],
            "avail": cols[3],
            "percent": cols[4],
            "mountpoint": cols[5],
        }
    return mp


def get_disks() -> List[Dict[str, Any]]:
    """
    시스템 디스크 및 파티션 정보를 lsblk 명령어로 가져오고,
    df 명령어로 사용량 정보를 매칭하며,
    smartctl로 온도 및 SMART 상태 정보를 실시간으로 조회해 리스트 반환
    """
    df_map = _df_map()

    rc, out = _run(["lsblk", "-J", "-o", "NAME,SIZE,FSTYPE,MOUNTPOINT,LABEL,TYPE"], timeout=6)
    if rc != 0 or not out.strip():
        return []

    try:
        data = json.loads(out)
    except Exception:
        return []

    disks: List[Dict[str, Any]] = []
    for dev in (data.get("blockdevices") or []):
        if not isinstance(dev, dict):
            continue

        name = str(dev.get("name") or "")
        dtype = str(dev.get("type") or "")
        if dtype != "disk":
            continue

        # 정책: nvme/sd만 노출
        if not (name.startswith("nvme") or name.startswith("sd")):
            continue

        d_df = df_map.get(name, {})
        dev_path = f"/dev/{name}"

        temp = get_disk_temperature(dev_path)
        smart_status = get_disk_smart_status(dev_path)

        parts: List[Dict[str, Any]] = []
        rep_mount = ""
        for part in (dev.get("children") or []):
            if not isinstance(part, dict):
                continue
            pname = str(part.get("name") or "")
            p_df = df_map.get(pname, {})
            pm = (part.get("mountpoint") or p_df.get("mountpoint") or "").strip()
            if pm and not rep_mount:
                rep_mount = pm

            parts.append({
                "name": pname,
                "size": str(part.get("size") or ""),
                "fstype": str(part.get("fstype") or ""),
                "mountpoint": pm,
                "label": str(part.get("label") or ""),
                "type": str(part.get("type") or ""),
                "used": p_df.get("used", ""),
                "avail": p_df.get("avail", ""),
                "percent": p_df.get("percent", ""),
            })

        disks.append({
            "name": name,
            "size": str(dev.get("size") or ""),
            "fstype": str(dev.get("fstype") or ""),
            "mountpoint": (dev.get("mountpoint") or d_df.get("mountpoint") or rep_mount or ""),
            "label": str(dev.get("label") or ""),
            "type": dtype,
            "used": d_df.get("used", ""),
            "avail": d_df.get("avail", ""),
            "percent": d_df.get("percent", ""),
            "temperature": temp,
            "smart_status": smart_status,
            "parts": parts,
        })

    return disks
