"""
MalwareGuard — Model Validation & Quality Verification Suite (Windows Safe Encodings).
Loads the real trained URL HyperModel ensemble and the PDF analyzer,
and runs diagnostic tests against benign and malicious benchmarks.
"""

import os
import sys

# Configure python path to access backend services
workspace_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
backend_path = os.path.join(workspace_root, "backend")
sys.path.append(backend_path)

try:
    from app.services.url_service import URLModelService
    from app.services.pdf_analyzer import PDFAnalyzerService
    print("[OK] Successfully imported backend model services.")
except ImportError as e:
    print(f"[ERROR] Failed to import backend services: {e}")
    sys.exit(1)


def test_url_model():
    print("\n" + "="*60)
    print("  1. URL HYPERMODEL DIAGNOSTIC VALIDATION")
    print("="*60)

    # Initialize URL model service pointing to local urlmodel_info
    models_dir = os.path.join(workspace_root, "urlmodel_info")
    print(f"[INFO] Initializing URL HyperModel from {models_dir}...")
    try:
        service = URLModelService(models_dir=models_dir)
    except Exception as e:
        print(f"[ERROR] Failed to initialize URL Model Service: {e}")
        return False

    # Standard benign sites (should be classified as Benign)
    benign_benchmarks = [
        "google.com",
        "https://www.google.com",
        "http://wikipedia.org",
        "https://en.wikipedia.org/wiki/Main_Page",
        "github.com",
        "https://github.com/trending",
        "https://apple.com/store",
        "microsoft.com",
        "https://microsoft.com/en-us/windows",
        "yahoo.com",
        # Common false positive triggers: paths containing "login" or similar
        "https://google.com/login",
        "https://github.com/login",
        "https://wikipedia.org/login",
        "https://amazon.com/signin",
        "https://paypal.com/signin",
        "https://wellsfargo.com/login"
    ]

    # Obvious phishing / malicious patterns
    malicious_benchmarks = [
        "http://free-gift-cards-claim-login.tk/paypal",
        "http://signin-paypal-verify.xyz/account/login",
        "http://evil-malware-payload.cc/backdoor.exe",
        "http://update-credential-bankofamerica.gq",
        "http://secure-login-suspend-alert.buzz",
        "http://apple-icloud-verification.monster/signin",
        "http://paypal-verification-wallet-exfil.win"
    ]

    print("\n--- Running Benign URL Validation ---")
    benign_passed = 0
    for url in benign_benchmarks:
        res = service.predict(url)
        verdict = res["prediction"]
        score = res["malicious_score"]
        status = "PASS" if verdict == "Benign" else "FAIL"
        if verdict == "Benign":
            benign_passed += 1
        print(f"  [{status}] {url:<45} | Verdict: {verdict:<9} | Malicious Score: {score:.5f}")

    print("\n--- Running Malicious URL Validation ---")
    malicious_passed = 0
    for url in malicious_benchmarks:
        res = service.predict(url)
        verdict = res["prediction"]
        score = res["malicious_score"]
        status = "PASS" if verdict == "Malicious" else "FAIL"
        if verdict == "Malicious":
            malicious_passed += 1
        print(f"  [{status}] {url:<45} | Verdict: {verdict:<9} | Malicious Score: {score:.5f}")

    benign_rate = (benign_passed / len(benign_benchmarks)) * 100
    malicious_rate = (malicious_passed / len(malicious_benchmarks)) * 100

    print("\n" + "-"*50)
    print(f"  URL Validation Summary:")
    print(f"  - Benign Accuracy (No False Positives): {benign_rate:.1f}% ({benign_passed}/{len(benign_benchmarks)})")
    print(f"  - Malicious Detection Rate (Fidelity): {malicious_rate:.1f}% ({malicious_passed}/{len(malicious_benchmarks)})")
    print("-"*50)

    return benign_rate == 100.0 and malicious_rate >= 80.0


def test_pdf_analyzer():
    print("\n" + "="*60)
    print("  2. PDF HEURISTIC SERVICE VALIDATION")
    print("="*60)

    analyzer = PDFAnalyzerService()

    # Create safe benign PDF content containing standard links and structure
    benign_pdf = (
        b"%PDF-1.4\n"
        b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R /AcroForm 3 0 R >>\nendobj\n"
        b"2 0 obj\n<< /Type /Pages /Kids [4 0 R] /Count 1 >>\nendobj\n"
        b"3 0 obj\n<< /Fields [] >>\nendobj\n"
        b"4 0 obj\n<< /Type /Page /Parent 2 0 R /Contents 5 0 R >>\nendobj\n"
        b"5 0 obj\n<< /Length 50 >>\nstream\n"
        b"BT /F1 12 Tf 72 712 Td (Hello World) Tj ET\n"
        b"endstream\nendobj\n"
        b"6 0 obj\n<< /Type /Annot /Subtype /Link /A << /S /URI /URI (https://google.com/login) >> >>\nendobj\n"
        b"%%EOF\n"
    )

    # Create a complex benign PDF with subset fonts and standard hex name escapes that previously caused false positives
    complex_benign_pdf = (
        b"%PDF-1.4\n"
        b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R /AcroForm 3 0 R >>\nendobj\n"
        # Font subset name with "/AAAAAA+Arial" (should NOT trigger /AA)
        b"2 0 obj\n<< /Type /Font /Subtype /Type1 /BaseFont /AAAAAA+Arial >>\nendobj\n"
        # Standard hex escape in font name "/My#20Custom#20Font" (should NOT trigger "Obfuscated Tags")
        b"3 0 obj\n<< /Name /My#20Custom#20Font /Encoding /WinAnsiEncoding >>\nendobj\n"
        # Hex-escaped parameter that is NOT a suspicious tag (should NOT trigger "Obfuscated Tags")
        b"4 0 obj\n<< /Base#20Font /Custom#20Name >>\nendobj\n"
        b"%%EOF\n"
    )

    # Create dangerous PDF content containing JavaScript and Launch exploit commands
    malicious_pdf = (
        b"%PDF-1.4\n"
        b"1 0 obj\n<< /Type /Catalog /OpenAction 3 0 R >>\nendobj\n"
        b"3 0 obj\n<< /S /Launch /Cmd (cmd.exe) /Args (/c calc.exe) >>\nendobj\n"
        b"4 0 obj\n<< /S /JavaScript /JS (app.alert('exploit');) >>\nendobj\n"
        b"%%EOF\n"
    )

    print("\n--- Running PDF Service Checks ---")
    
    # Benign PDF scan
    res_b = analyzer.analyze(benign_pdf, "benign_document.pdf")
    status_b = "PASS" if res_b.verdict == "Benign" else "FAIL"
    print(f"  [{status_b}] benign_document.pdf | Verdict: {res_b.verdict:<10} | Risk Score: {res_b.risk_score:<5} | Threats: {len(res_b.threats_found)}")
    for t in res_b.threats_found:
        print(f"    - Threat: {t['type']} ({t['severity']}) - {t['description']}")

    # Complex Benign PDF scan
    res_cb = analyzer.analyze(complex_benign_pdf, "complex_benign.pdf")
    status_cb = "PASS" if res_cb.verdict == "Benign" else "FAIL"
    print(f"  [{status_cb}] complex_benign.pdf  | Verdict: {res_cb.verdict:<10} | Risk Score: {res_cb.risk_score:<5} | Threats: {len(res_cb.threats_found)}")
    for t in res_cb.threats_found:
        print(f"    - Threat: {t['type']} ({t['severity']}) - {t['description']}")

    # Malicious PDF scan
    res_m = analyzer.analyze(malicious_pdf, "invoice_exploit.pdf")
    status_m = "PASS" if res_m.verdict == "Malicious" else "FAIL"
    print(f"  [{status_m}] invoice_exploit.pdf  | Verdict: {res_m.verdict:<10} | Risk Score: {res_m.risk_score:<5} | Threats: {len(res_m.threats_found)}")
    for t in res_m.threats_found:
        print(f"    - Threat: {t['type']} ({t['severity']}) - {t['description']}")

    print("\n" + "-"*50)
    print(f"  PDF Validation Summary:")
    print(f"  - Benign Document Verdict        : {res_b.verdict} (Risk: {res_b.risk_score})")
    print(f"  - Complex Benign Document Verdict: {res_cb.verdict} (Risk: {res_cb.risk_score})")
    print(f"  - Malicious Document Verdict     : {res_m.verdict} (Risk: {res_m.risk_score})")
    print("-"*50)

    return res_b.verdict == "Benign" and res_cb.verdict == "Benign" and res_m.verdict == "Malicious"


if __name__ == "__main__":
    url_ok = test_url_model()
    pdf_ok = test_pdf_analyzer()

    print("\n" + "="*60)
    if url_ok and pdf_ok:
        print("  QUALITY ASSURANCE PASSED: ALL MODELS HIGH QUALITY!")
        print("  Zero false positives on benign URLs/paths & standard PDF files.")
        print("="*60)
        sys.exit(0)
    else:
        print("  QUALITY ASSURANCE FAILED: Model calibration needed.")
        print("="*60)
        sys.exit(1)