"""
generate_aibom_pdf.py — Generate AIBoM + PDF Technical File for medassist-tinyllama.

Reads:  results_medassist-tinyllama.json  (produced by test_local_medassist.py)
Writes:
  aibom_medassist-tinyllama.json   — CycloneDX 1.6 ML-BOM
  technical_file_medassist-tinyllama.pdf  — EU AI Act Annex IV Technical File

Usage:
    PYTHONPATH="$PWD" python3 sichgate/generate_aibom_pdf.py
"""

import json
import os
import sys
from pathlib import Path

# ── Paths ─────────────────────────────────────────────────────────────────────
RESULTS_JSON   = "results_medassist-tinyllama.json"
ADAPTER_PATH   = (
    "/Users/polinamoshenets/Documents/healthcare-model/"
    "Custom Med-Model/medassist-tinyllama"
)
AIBOM_OUT      = "aibom_medassist-tinyllama.json"
PDF_OUT        = "technical_file_medassist-tinyllama.pdf"
MODEL_LABEL    = "medassist-tinyllama"
MODEL_VERSION  = "1.0.0-lora-tinyllama"
# ─────────────────────────────────────────────────────────────────────────────

from sichgate.aibom_capture import AIBoMCapture, capture_model_aibom
from sichgate.technical_file_builder import TechnicalFileBuilder
from sichgate.pdf_exporter import PDFExporter

# ── 1. Load findings from the audit report ────────────────────────────────────
print(f"Loading findings from {RESULTS_JSON} ...")
with open(RESULTS_JSON, encoding="utf-8") as f:
    report = json.load(f)

# The report stores results as a list of attack result dicts
findings: list = report.get("results", [])
print(f"  [{chr(10003)}] {len(findings)} findings loaded")

# ── 2. Build AIBoM ────────────────────────────────────────────────────────────
# Point model_path at the adapter weights file for SHA-256 integrity hashing
adapter_weights = os.path.join(ADAPTER_PATH, "adapter_model.safetensors")

print(f"\nCapturing AIBoM for {MODEL_LABEL} ...")
aibom_capture = AIBoMCapture(
    model_path=adapter_weights,
    model_name=MODEL_LABEL,
    model_version=MODEL_VERSION,
    quantization_type=None,   # full precision LoRA adapter (not quantized)
)
cyclonedx_bom = aibom_capture.to_cyclonedx()

# Embed base model info as additional properties
cyclonedx_bom["metadata"]["properties"].append(
    {"name": "adapter:base_model",      "value": "TinyLlama/TinyLlama-1.1B-Chat-v1.0"}
)
cyclonedx_bom["metadata"]["properties"].append(
    {"name": "adapter:peft_type",       "value": "LORA"}
)
cyclonedx_bom["metadata"]["properties"].append(
    {"name": "adapter:lora_rank",       "value": "16"}
)
cyclonedx_bom["metadata"]["properties"].append(
    {"name": "adapter:lora_alpha",      "value": "32"}
)
cyclonedx_bom["metadata"]["properties"].append(
    {"name": "adapter:target_modules",  "value": "q_proj,v_proj"}
)
cyclonedx_bom["metadata"]["properties"].append(
    {"name": "adapter:task_type",       "value": "CAUSAL_LM"}
)
cyclonedx_bom["metadata"]["properties"].append(
    {"name": "adapter:local_path",      "value": ADAPTER_PATH}
)

# Save AIBoM JSON
aibom_path = Path(AIBOM_OUT)
aibom_path.write_text(json.dumps(cyclonedx_bom, indent=2, ensure_ascii=False))
print(f"  [{chr(10003)}] AIBoM saved  → {AIBOM_OUT}")
print(f"        SHA-256 of adapter weights: "
      f"{aibom_capture.model_sha256 or '(file not found)'}")

# ── 3. Build EU AI Act Annex IV Technical File ────────────────────────────────
print(f"\nBuilding Annex IV Technical File ...")
builder = TechnicalFileBuilder(
    findings=findings,
    aibom=cyclonedx_bom,
)
technical_file = builder.build()
print(f"  [{chr(10003)}] Technical File built — {len(technical_file.sections)} sections")

# ── 4. Export PDF ─────────────────────────────────────────────────────────────
print(f"\nExporting PDF ...")
exporter = PDFExporter()
try:
    pdf_path = exporter.export_pdf(technical_file, PDF_OUT)
    print(f"  [{chr(10003)}] PDF saved → {pdf_path}")
except Exception as e:
    # Fallback: save HTML if PDF renderer not available
    html_out = PDF_OUT.replace(".pdf", ".html")
    html_content = exporter.render_html(technical_file)
    Path(html_out).write_text(html_content, encoding="utf-8")
    print(f"  [!] PDF renderer unavailable ({e})")
    print(f"  [{chr(10003)}] HTML Technical File saved → {html_out}")
    print(f"      Install weasyprint for PDF:  pip install weasyprint")

print(f"\nDone.")
print(f"  AIBoM:          {AIBOM_OUT}")
print(f"  Technical File: {PDF_OUT}")
