### test_local_medassist.py — tests the locally trained medassist-tinyllama model ###
"""
SichGate Pro test harness for a locally trained LoRA model.

Model: medassist-tinyllama
  - Base:    TinyLlama/TinyLlama-1.1B-Chat-v1.0
  - Adapter: LoRA (CAUSAL_LM), rank=16, alpha=32, targets=[q_proj, v_proj]
  - Path:    /Users/polinamoshenets/Documents/healthcare-model/
             Custom Med-Model/medassist-tinyllama

Loading strategy:
  1. AutoTokenizer from adapter path (picks up custom chat_template.jinja)
  2. AutoModelForCausalLM for base TinyLlama (CPU, avoids MPS buffer limits)
  3. PeftModel applies the LoRA weights, then merge_and_unload() for clean inference
  4. text-generation pipeline wraps the merged model

Output: results_medassist-tinyllama.json  (auditor-ready SichGateReport)
"""

import json
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from peft import PeftModel

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

# ── Model paths ───────────────────────────────────────────────────────────────
BASE_MODEL_ID   = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
ADAPTER_PATH    = (
    "/Users/polinamoshenets/Documents/healthcare-model/"
    "Custom Med-Model/medassist-tinyllama"
)
MODEL_LABEL     = "medassist-tinyllama"   # used in output filename / report
# ─────────────────────────────────────────────────────────────────────────────

print(f"Loading base model: {BASE_MODEL_ID} ...")
# Load tokenizer from the base model — the local tokenizer_config uses
# 'TokenizersBackend' which is not a standard Transformers class.
# We then overlay the custom chat template from the adapter directory.
tokenizer = AutoTokenizer.from_pretrained(
    BASE_MODEL_ID,
    trust_remote_code=False,
)
# Apply the fine-tuned chat template if present
_chat_template_path = os.path.join(ADAPTER_PATH, "chat_template.jinja")
if os.path.exists(_chat_template_path):
    with open(_chat_template_path, encoding="utf-8") as _f:
        tokenizer.chat_template = _f.read()
    print(f"  [✓] Custom chat template loaded from adapter")

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    dtype=torch.float32,         # float32 for CPU stability
    device_map="cpu",
    trust_remote_code=False,
)

print(f"Applying LoRA adapter from: {ADAPTER_PATH} ...")
peft_model = PeftModel.from_pretrained(base_model, ADAPTER_PATH)
merged_model = peft_model.merge_and_unload()   # fuse LoRA weights for clean inference
merged_model.eval()

print("Building text-generation pipeline ...")
gen_pipe = pipeline(
    "text-generation",
    model=merged_model,
    tokenizer=tokenizer,
    device_map="cpu",
)


def model_fn(prompt: str) -> str:
    """
    Black-box generative adapter.

    Wraps the prompt in the model's chat template (<|user|> / <|assistant|>),
    generates up to 256 new tokens, and returns the assistant's reply only.

    Returns '[No output]' when the model produces an empty continuation.
    """
    try:
        messages = [{"role": "user", "content": prompt[:1024]}]
        formatted = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
        outputs = gen_pipe(
            formatted,
            max_new_tokens=96,            # keep short for faster CPU adversarial testing
            do_sample=False,              # deterministic for reproducible audit
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
            return_full_text=False,       # return assistant reply only
        )
        text = outputs[0]["generated_text"].strip()
        return text if text else "[No output]"
    except Exception as exc:
        return f"[Generation error: {exc}]"


# ── Regulatory context ────────────────────────────────────────────────────────
# Healthcare assistant in a clinical/pharma context:
#   - Annex III point 5 (healthcare AI) → annex_iii_applicable = True
#   - HIPAA applies (processes patient-facing health queries)
#   - bias_threshold = 0.15 (healthcare, stricter than 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,
    key_deadlines={
        "EU_AI_Act_high_risk": "2026-08-02",
        "GDPR": "in_force",
        "HIPAA": "in_force",
        "NIST_AI_RMF": "voluntary",
        "ISO_42001": "certification_available",
    },
)

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

config = SichGateConfig(
    model_fn=model_fn,
    model_target=MODEL_LABEL,

    # 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,         # Role-play / persona boundary testing
    run_json_xml_smuggling=True,     # OWASP LLM01
    run_textattack=True,             # Gracefully skips if textattack not installed

    # ── Customer context ──────────────────────────────────────────────────────
    hipaa_mode=True,                 # Adds HIPAA §164.308 / §164.312 compliance tags
    biometric_mode=False,            # Not a biometric identification model

    # ── 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 ────────────────────────────────────────────────
audit_report = AuditReport.from_attack_results(
    results=runner_report.results,
    model_target=MODEL_LABEL,
    report_id=runner_report.report_id,
    license_key_hash="DEMO",          # Replace with real key hash in production
    model_version=None,               # Pin to a commit SHA in production
    parameters_millions=1100,         # TinyLlama-1.1B ≈ 1100 M params
    regulatory_context=regulatory_context,
)

output_path = f"results_{safe_name}.json"
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=ADAPTER_PATH,
    model_name=MODEL_LABEL,
    model_version="local-lora-v1",
)
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 vars(r)
    for r in audit_report.results
]
print("  Building Technical File...", flush=True)
technical_file = TechnicalFileBuilder(
    findings=findings_list,
    aibom=aibom_dict,
).build()

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"
    PDFExporter().render_html(technical_file)
    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 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"      frameworks: {', '.join(regulatory_context.applicable_frameworks)}")
