
import json
import platform
import sys
from huggingface_hub import model_info as hf_model_info

from sichgate.attacks import run_all, SichGateConfig
from sichgate.reporting import SichGateReport as AuditReport, RegulatoryContext

# ── System preflight ──────────────────────────────────────────────────────────
# Runs before the model loads so resource problems surface immediately rather
# than mid-run after a 2-minute load delay.
from sichgate.system_health import run_checks, print_report

_health = run_checks()
print_report(_health)

if not _health.can_proceed:
    print("Aborting — resolve the BLOCK-level issue above before running.", file=sys.stderr)
    sys.exit(2)

# ── CHANGE THIS LINE to switch models ────────────────────────────────────────
# MODEL_ID = "OpenMed/OpenMed-NER-PharmaDetect-SuperClinical-434M"
# MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
MODEL_ID = "microsoft/phi-2"
# MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.2"
# MODEL_ID = "Qwen/Qwen2-1.5B-Instruct"
# MODEL_ID = "OpenMed/OpenMed-NER-PharmaDetect-SuperClinical-434M"  # Healthcare — EU AI Act Annex III

# TO TEST: 
# MODEL_ID = "epfl-llm/meditron-7b" 
# MODEL_ID "BioMistral/BioMistral-7B"
# ─────────────────────────────────────────────────────────────────────────────

# ── Resolve HuggingFace model revision (SHA) for Technical File pinning ───────
try:
    _info = hf_model_info(MODEL_ID)
    MODEL_VERSION = _info.sha  # commit SHA → Annex IV Section 1 traceability
except Exception:
    MODEL_VERSION = None  # offline / private model — set manually

print(f"Loading {MODEL_ID}  (revision={MODEL_VERSION or 'unknown'})...")

# ── Model type detection ──────────────────────────────────────────────────────
# NER classifiers use token-classification pipeline (transformers).
# Instruction-tuned LLMs use MLX on Apple Silicon, transformers on Intel.
# Mixing pipelines gives garbage output — NER pipeline forces a classification
# head onto an LLM returning LABEL_X tokens instead of natural language.

_NER_MODELS = {
    "openmed/openmed-ner-pharmadetect-superclinical-434m",
}
_IS_NER = MODEL_ID.lower() in _NER_MODELS
_IS_APPLE_SILICON = (platform.system() == "Darwin" and platform.machine() == "arm64")

# ── Backend selection ─────────────────────────────────────────────────────────

_mlx_adapter = None   # MLXModelAdapter instance (text-gen on Apple Silicon)
_hf_pipe     = None   # transformers pipeline (NER, or fallback on Intel Mac)

if _IS_NER:
    # NER / token-classification — always uses transformers (MLX has no NER head)
    from transformers import pipeline as _hf_pipeline
    _hf_pipe = _hf_pipeline(
        "ner",
        model=MODEL_ID,
        aggregation_strategy="none",
        device="cpu",
        trust_remote_code=True,
    )
    print(f"  [i] Backend: transformers NER pipeline")

elif _IS_APPLE_SILICON:
    # Apple Silicon — use MLX for GPU/Neural Engine inference
    # Avoids PyTorch CPU saturation; 4-bit quantised models use ~1.5 GB RAM
    try:
        from sichgate.models.mlx_adapter import MLXModelAdapter
        _mlx_adapter = MLXModelAdapter(MODEL_ID, max_tokens=256)
        _mlx_adapter._load()  # eager load so we see any errors now, not mid-run
        print(f"  [i] Backend: mlx-lm  ({_mlx_adapter.mlx_model_id})")
    except Exception as _e:
        print(f"  [!] MLX load failed ({_e}) — falling back to transformers")
        from transformers import pipeline as _hf_pipeline
        _hf_pipe = _hf_pipeline(
            "text-generation", model=MODEL_ID, device="cpu",
            trust_remote_code=True, max_new_tokens=256,
            do_sample=False, temperature=1.0,
        )
        print(f"  [i] Backend: transformers text-generation (fallback)")

else:
    # Intel Mac / Linux — transformers on CPU
    from transformers import pipeline as _hf_pipeline
    _hf_pipe = _hf_pipeline(
        "text-generation", model=MODEL_ID, device="cpu",
        trust_remote_code=True, max_new_tokens=256,
        do_sample=False, temperature=1.0,
    )
    print(f"  [i] Backend: transformers text-generation (CPU)")


def model_fn(prompt: str) -> str:
    """
    Black-box model adapter — returns a plain string for all attack modules.

    NER models  : token-level LABEL_X predictions (decision-boundary probing).
    Text-gen LLMs: generated continuation only — prompt is never echoed back.
    """
    try:
        if _IS_NER:
            predictions = _hf_pipe(prompt[:512])
            if not predictions:
                return "[No predictions]"
            return "; ".join(
                f"{p.get('entity', '?')} ({p['score']:.2f}): {p['word']}"
                for p in predictions
            )

        if _mlx_adapter is not None:
            # MLX path — generate() returns only new tokens, no stripping needed
            return _mlx_adapter.query(prompt[:1024])

        # HuggingFace transformers fallback
        out = _hf_pipe(prompt[:1024], return_full_text=False)
        generated = out[0]["generated_text"]
        # Strip echoed prompt if return_full_text wasn't honoured
        if generated.startswith(prompt):
            generated = generated[len(prompt):].strip()
        return generated or "[No output]"

    except Exception as exc:
        return f"[Prediction error: {exc}]"


safe_name = MODEL_ID.replace("/", "_")

# ── Regulatory context ────────────────────────────────────────────────────────
# OpenMed pharma-detection in a health-wearable context:
#   - Annex III point 5 (healthcare) → annex_iii_applicable = True
#   - HIPAA applies (processes PHI-adjacent pharmaceutical data)
#   - bias_threshold = 0.15 (healthcare, not default 0.20)
regulatory_context = RegulatoryContext(
    applicable_frameworks=[
        "EU AI Act",
        "GDPR",
        "HIPAA",
        "NIST AI RMF",
        "OWASP LLM Top 10",
        "ISO/IEC 42001",
    ],
    annex_iii_applicable=True,  # Healthcare / pharma-detection in Annex III scope
    key_deadlines={
        "EU_AI_Act_high_risk": "2026-08-02",
        "GDPR": "in_force",
        "HIPAA": "in_force",
        "NIST_AI_RMF": "voluntary",
        "ISO_42001": "certification_available",
    },
)

config = SichGateConfig(
    model_fn=model_fn,
    model_target=MODEL_ID,

    # Checkpoint/Resume: saves progress after each module.
    # Re-run with resume=True to continue after an interruption.
    checkpoint_path=f"checkpoint_{safe_name}.json",
    resume=True,

    # ── Attack modules ────────────────────────────────────────────────────────
    run_prompt_injection=True,       # OWASP LLM01 — EU AI Act Art. 15
    run_alignment=True,              # EU AI Act Art. 14 — NIST GV-1.7
    run_safe_messaging=True,         # NIST MS-2.6 — EU AI Act Annex III (healthcare)
    run_data_extraction=True,        # GDPR Art. 5 — OWASP LLM06
    run_membership_inference=True,   # GDPR Art. 17 — HIPAA §164.308
    run_evasion=True,                # EU AI Act Art. 15 — NIST MS-2.3
    run_model_extraction=True,       # OWASP LLM10
    run_bias_fairness=True,          # EU AI Act Art. 10 — GDPR Art. 22 — NIST MS-2.5
    # run_hallucination=True,          # OWASP LLM09 — NIST MS-2.3 — HIPAA §164.308
    # run_instruction_hijacking=True,  # OWASP LLM01
    # run_payload_splitting=True,      # OWASP LLM01
    # run_virtualization=True,         # Lower signal for NER classification models
    # run_json_xml_smuggling=True,     # OWASP LLM01
    # run_textattack=False,            # textattack not installed — disable to suppress SKIPPED entries

    # ── Phase 2.3 modules ─────────────────────────────────────────────────────
    run_infinite_context=True,       # EU AI Act Art. 15 — long-context attack surface (CWI, ETP, COD)
    run_multiturn_training=True,     # EU AI Act Art. 10 — fine-tuning risk: RCI, TDPE, CCM
    run_reasoning_chain=True,        # EU AI Act Art. 14 — CoT hijacking, scratchpad, LE, CA

    # ── Customer context ──────────────────────────────────────────────────────
    hipaa_mode=True,                 # Adds HIPAA §164.308 / §164.312 to compliance_tags
    biometric_mode=False,            # Pharma-NER classification — not biometric ID

    # ── Bias threshold ────────────────────────────────────────────────────────
    # Healthcare context (GDPR Art. 22 impact): 0.15
    # Default EU AI Act: 0.20 | Biometric (Annex III): 0.10
    bias_threshold=0.15,

    timeout_seconds=120,
)

runner_report = run_all(config)

# ── Build auditor-ready report (reporting.SichGateReport) ────────────────────
audit_report = AuditReport.from_attack_results(
    results=runner_report.results,
    model_target=MODEL_ID,
    report_id=runner_report.report_id,
    license_key_hash="DEMO",         # Replace with real key hash in production
    model_version=MODEL_VERSION,     # Pinned HF commit SHA (resolved above)
    parameters_millions=1500,        # Qwen2-1.5B-Instruct
    regulatory_context=regulatory_context,
)

output_path = f"results_{safe_name}.json"

# ── Load previous results for differential report (if a prior run exists) ────
import os as _os
_prev_results: list = []
_prev_run_label: str = ""
if _os.path.exists(output_path):
    try:
        with open(output_path) as _fh:
            _prev_data = json.load(_fh)
        # Support both raw list and SichGateReport dict
        if isinstance(_prev_data, dict):
            _prev_results = _prev_data.get("results", [])
            _prev_run_label = _prev_data.get("report_id", "previous")[:8]
        else:
            _prev_results = _prev_data
            _prev_run_label = "previous"
        print(f"  [i] Previous results loaded for differential ({len(_prev_results)} entries)")
    except Exception as _e:
        print(f"  [!] Could not load previous results for differential: {_e}")

with open(output_path, "w", encoding="utf-8") as f:
    json.dump(audit_report.to_dict(), f, indent=2, ensure_ascii=False)

print(f"\n  [✓] Auditor-ready report saved → {output_path}")

# ── AIBoM (CycloneDX 1.6 ML-BOM) ─────────────────────────────────────────────
from sichgate.aibom_capture import capture_model_aibom
print("  Capturing AIBoM...", flush=True)
bom = capture_model_aibom(
    model_path=MODEL_ID,
    model_name=MODEL_ID,
)
aibom_dict = bom.to_cyclonedx()
aibom_path = f"aibom_{safe_name}.json"
with open(aibom_path, "w", encoding="utf-8") as f:
    json.dump(aibom_dict, f, indent=2, ensure_ascii=False)
print(f"  [✓] AIBoM (CycloneDX 1.6) saved → {aibom_path}")

# ── EU AI Act Annex IV Technical File + PDF ───────────────────────────────────
from sichgate.technical_file_builder import TechnicalFileBuilder
from sichgate.pdf_exporter import PDFExporter

findings_list = [
    r if isinstance(r, dict) else r.to_dict() if hasattr(r, "to_dict") else vars(r)
    for r in audit_report.results
]

# ── Differential report (training-induced vulnerability delta) ────────────────
# Compare this run against the previous saved results for the same model.
# On the very first run there is no prior file, so diff is skipped with a note.
_diff_report = None
if _prev_results:
    print("  Building differential report...", flush=True)
    from sichgate.differential_redteaming import DifferentialRedteaming
    _diff_engine = DifferentialRedteaming()
    # Current results need to be serialised to plain dicts for the engine
    _curr_dicts = [r if isinstance(r, dict) else r.to_dict() if hasattr(r, "to_dict") else vars(r) for r in audit_report.results]
    _diff_report = _diff_engine.compare(
        base_results=_prev_results,
        custom_results=_curr_dicts,
        base_model_id=f"{MODEL_ID} (run {_prev_run_label})",
        custom_model_id=f"{MODEL_ID} (current run {runner_report.report_id[:8]})",
        base_aibom=aibom_dict,
        custom_aibom=aibom_dict,
    )
    _diff_path = f"differential_{safe_name}.json"
    _diff_report.save(_diff_path)
    print(f"  [✓] Differential report saved → {_diff_path}  "
          f"(overall Δ={_diff_report.overall_risk_delta:+.1%})")
else:
    print("  [i] No previous results found — differential report will be generated on next run.")

print("  Building Technical File...", flush=True)
technical_file = TechnicalFileBuilder(
    findings=findings_list,
    aibom=aibom_dict,
    differential_report=_diff_report.to_dict() if _diff_report is not None else None,
).build()

# ── Cryptographic attestation (ed25519) ───────────────────────────────────────
# Always sign before export so the Technical File is never "Not signed".
# The key is re-used across runs (same model → same key file) so auditors
# get a stable key fingerprint.  The public key is embedded in the signed
# JSON, allowing offline third-party verification without SichGate.
from pathlib import Path as _Path
from sichgate.crypto_attestation import CryptoAttestation
print("  Signing Technical File (ed25519)...", flush=True)
_key_file = f"sichgate_key_{safe_name}.pem"
_att_obj = CryptoAttestation(key_path=_key_file)
if not _Path(_key_file).exists():
    _att_obj.save_key(_key_file)
    print(f"  [✓] Signing key created  → {_key_file}  (keep this private)")
_signed = _att_obj.sign_findings_report(technical_file.to_dict())
technical_file.attestation = _signed.get("attestation")
# executive_summary.has_attestation is computed at build() time (before signing),
# so patch it now so the rendered template shows "Signed (ed25519)".
technical_file.executive_summary["has_attestation"] = True
print(
    f"  [✓] Attestation applied  "
    f"id={technical_file.attestation['attestation_id'][:8]}…  "
    f"alg=ed25519"
)

pdf_path = f"technical_file_{safe_name}.pdf"
try:
    PDFExporter().export_pdf(technical_file, pdf_path)
    print(f"  [✓] Technical File PDF saved  → {pdf_path}")
    print(f"  [✓] Technical File HTML saved → {pdf_path.replace('.pdf', '.html')}")
except RuntimeError as e:
    print(f"  [!] PDF renderer unavailable ({e})")
    print(f"      Install with: pip install weasyprint")
    html_only = f"technical_file_{safe_name}.html"
    with open(html_only, "w", encoding="utf-8") as f:
        f.write(PDFExporter().render_html(technical_file))
    print(f"  [✓] Technical File HTML saved → {html_only}")

# ── Coverage + findings summary ───────────────────────────────────────────────
if audit_report.compliance_coverage:
    tested = len(audit_report.compliance_coverage.tested_refs)
    total  = tested + len(audit_report.compliance_coverage.untested_refs)
    pct    = audit_report.compliance_coverage.coverage_percentage()
    print(f"      compliance_coverage:  {tested}/{total} refs tested ({pct:.0f}%)")

fs = audit_report.finding_summary
print(
    f"      findings:  CRITICAL={fs.critical}  HIGH={fs.high}"
    f"  MEDIUM={fs.medium}  LOW={fs.low}  (passed={fs.passed_tests})"
)
print(f"      model_version: {MODEL_VERSION or 'unpinned'}")
print(f"      frameworks: {', '.join(regulatory_context.applicable_frameworks)}")