#  Sovereign AIBoM capture in CycloneDX 1.6 ML-BOM format.

# Captures a cryptographically anchored snapshot of the model under test at
# the moment the scan begins.  The AIBoM constitutes provenance evidence for
# the Technical File required under EU AI Act Article 11.

# Output format: CycloneDX 1.6 JSON with ML-BOM extensions.

# Spec refs:
#   CycloneDX 1.6: https://cyclonedx.org/docs/1.6/json/
#   ML-BOM:        https://cyclonedx.org/capabilities/mlbom/
#   EU AI Act Art. 11 + Annex IV §1a, §1e, §2c, §2d

from __future__ import annotations

import hashlib
import json
import os
import platform
import subprocess
import uuid
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional


_SICHGATE_VERSION = "2.1.0"
_CYCLONEDX_SPEC = "1.6"
_CYCLONEDX_BOM_FORMAT = "CycloneDX"




@dataclass
class HardwareContext:
    """Hardware environment captured at scan time (Annex IV §1e)."""
    cpu_model: str
    cpu_cores: int
    ram_gb: float
    platform_os: str
    platform_arch: str
    gpu_model: Optional[str] = None
    gpu_vram_gb: Optional[float] = None

    def to_dict(self) -> dict:
        return {k: v for k, v in asdict(self).items() if v is not None}


def detect_hardware() -> HardwareContext:
    cpu_model = platform.processor() or "unknown"
    cpu_cores = os.cpu_count() or 1

    # RAM — best-effort across Linux / macOS
    ram_gb: float = 0.0
    try:
        import psutil  # type: ignore
        ram_gb = round(psutil.virtual_memory().total / (1024 ** 3), 1)
    except ImportError:
        try:
            if platform.system() == "Darwin":
                raw = subprocess.check_output(
                    ["sysctl", "-n", "hw.memsize"], text=True
                ).strip()
                ram_gb = round(int(raw) / (1024 ** 3), 1)
            elif platform.system() == "Linux":
                with open("/proc/meminfo") as fh:
                    for line in fh:
                        if line.startswith("MemTotal:"):
                            kb = int(line.split()[1])
                            ram_gb = round(kb / (1024 ** 2), 1)
                            break
        except Exception:
            pass

    # GPU — best-effort via nvidia-smi or PyTorch
    gpu_model: Optional[str] = None
    gpu_vram_gb: Optional[float] = None
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=name,memory.total",
             "--format=csv,noheader,nounits"],
            text=True, stderr=subprocess.DEVNULL,
        ).strip()
        if out:
            parts = out.split(",")
            gpu_model = parts[0].strip()
            gpu_vram_gb = round(float(parts[1].strip()) / 1024, 1)
    except Exception:
        try:
            import torch  # type: ignore
            if torch.cuda.is_available():
                gpu_model = torch.cuda.get_device_name(0)
                gpu_vram_gb = round(
                    torch.cuda.get_device_properties(0).total_memory / (1024 ** 3), 1
                )
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                gpu_model = "Apple Silicon MPS"
        except Exception:
            pass

    return HardwareContext(
        cpu_model=cpu_model,
        cpu_cores=cpu_cores,
        ram_gb=ram_gb,
        platform_os=f"{platform.system()} {platform.release()}",
        platform_arch=platform.machine(),
        gpu_model=gpu_model,
        gpu_vram_gb=gpu_vram_gb,
    )



def _sha256_file(path: str) -> Optional[str]:
    p = Path(path)
    if not p.is_file():
        return None
    h = hashlib.sha256()
    with p.open("rb") as fh:
        for chunk in iter(lambda: fh.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def _file_size(path: str) -> Optional[int]:
    p = Path(path)
    return p.stat().st_size if p.is_file() else None


def _infer_format(path: str) -> Optional[str]:
    ext = Path(path).suffix.lstrip(".").lower()
    return ext if ext else None


#aibom captue 

@dataclass
class AIBoMCapture:
    model_path: str
    model_version: str
    model_name: str
    quantization_type: Optional[str] = None
    test_run_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    hardware: HardwareContext = field(default_factory=detect_hardware)
    timestamp: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )

    # Computed from model_path during __post_init__
    model_sha256: Optional[str] = field(default=None, init=False)
    model_file_size_bytes: Optional[int] = field(default=None, init=False)
    model_format: Optional[str] = field(default=None, init=False)

    # EU AI Act Article 11 / Annex IV section evidence
    article_11_sections: List[str] = field(
        default_factory=lambda: ["1", "1a", "1e", "2c", "2d"]
    )

    def __post_init__(self) -> None:
        self.model_sha256 = _sha256_file(self.model_path)
        self.model_file_size_bytes = _file_size(self.model_path)
        self.model_format = _infer_format(self.model_path)

    # ── CycloneDX 1.6 serialisation ─────────────────────────────────────────

    def to_cyclonedx(self) -> Dict[str, Any]:
        
        # model componoont 
        component: Dict[str, Any] = {
            "type": "machine-learning-model",
            "bom-ref": f"model:{self.test_run_id}",
            "name": self.model_name,
            "version": self.model_version,
        }

        if self.model_sha256:
            component["hashes"] = [
                {"alg": "SHA-256", "content": self.model_sha256}
            ]

        # CycloneDX ML-BOM model card (ML-BOM extension)
        model_card: Dict[str, Any] = {
            "modelParameters": {
                "task": "text-generation",
                "architectureFamily": "transformer",
                "modelArchitecture": self.model_name,
                "datasets": [],
            },
            "quantitativeAnalysis": {},
        }
        if self.quantization_type:
            model_card["modelParameters"]["quantizationType"] = self.quantization_type
        component["modelCard"] = model_card

        # Component-level properties (file metadata)
        comp_props: List[Dict[str, str]] = [
            {"name": "sichgate:test_run_id", "value": self.test_run_id},
        ]
        if self.model_format:
            comp_props.append({"name": "sichgate:model_format", "value": self.model_format})
        if self.model_file_size_bytes is not None:
            comp_props.append(
                {"name": "sichgate:file_size_bytes",
                 "value": str(self.model_file_size_bytes)}
            )
        if self.quantization_type:
            comp_props.append(
                {"name": "sichgate:quantization_type", "value": self.quantization_type}
            )
        component["properties"] = comp_props

        #  Metadata hardware props
        hw = self.hardware
        hw_props: List[Dict[str, str]] = [
            {"name": "hardware:cpu_model",     "value": hw.cpu_model},
            {"name": "hardware:cpu_cores",     "value": str(hw.cpu_cores)},
            {"name": "hardware:ram_gb",        "value": str(hw.ram_gb)},
            {"name": "hardware:platform_os",   "value": hw.platform_os},
            {"name": "hardware:platform_arch", "value": hw.platform_arch},
        ]
        if hw.gpu_model:
            hw_props.append({"name": "hardware:gpu_model", "value": hw.gpu_model})
        if hw.gpu_vram_gb is not None:
            hw_props.append(
                {"name": "hardware:gpu_vram_gb", "value": str(hw.gpu_vram_gb)}
            )

        # Annex IV Article 11 mapping
        annex_props: List[Dict[str, str]] = [
            {
                "name": "eu_ai_act:article_11_sections",
                "value": ",".join(self.article_11_sections),
            },
        ]

        return {
            "bomFormat": _CYCLONEDX_BOM_FORMAT,
            "specVersion": _CYCLONEDX_SPEC,
            "version": 1,
            "serialNumber": f"urn:uuid:{self.test_run_id}",
            "metadata": {
                "timestamp": self.timestamp,
                "tools": [
                    {
                        "vendor": "SichGate",
                        "name": "SichGate Pro",
                        "version": _SICHGATE_VERSION,
                    }
                ],
                # metadata.component = the subject being described (the model under test),
                # NOT the scanning tool.  CycloneDX spec §3.3: "The component element …
                # describes the component the BOM is for."
                "component": {
                    "type": "machine-learning-model",
                    "name": self.model_name,
                    "version": self.model_version,
                },
                "properties": hw_props + annex_props,
            },
            "components": [component],
        }

    def to_dict(self) -> Dict[str, Any]:
        return self.to_cyclonedx()

    def to_json(self, indent: int = 2) -> str:
        return json.dumps(self.to_cyclonedx(), indent=indent)

    def save(self, path: str) -> None:
        Path(path).write_text(self.to_json())



def capture_model_aibom(
    model_path: str,
    model_name: Optional[str] = None,
    model_version: str = "unknown",
    quantization_type: Optional[str] = None,
) -> AIBoMCapture:
  
    name = model_name or Path(model_path).stem or model_path
    return AIBoMCapture(
        model_path=model_path,
        model_name=name,
        model_version=model_version,
        quantization_type=quantization_type,
    )
