#!/usr/bin/env python3
"""
╔═══════════════════════════════════════════════════════════════════════╗
║  🔱  ZKAEDI PRIME — UNIFIED SECURITY AUDIT PIPELINE                 ║
║                                                                       ║
║  Chains all 4 HuggingFace models/tools into a single audit command:  ║
║                                                                       ║
║  Stage 1: gemma-2-9b-solidity-merged                                 ║
║           → Generates vulnerability energy signatures from Solidity   ║
║                                                                       ║
║  Stage 2: prime-swarm-hunter (Gradio Space)                          ║
║           → 12-agent temporal compound detection                      ║
║                                                                       ║
║  Stage 3: leviathan-v2                                               ║
║           → CNN exploit topology classification + PRIME refinement    ║
║                                                                       ║
║  Stage 4: solidity-vuln-auditor-7b                                   ║
║           → Synthesizes final professional audit report               ║
║                                                                       ║
║  Author: ZKAEDI — Offensive Healer                                    ║
╚═══════════════════════════════════════════════════════════════════════╝

USAGE:
    # Full pipeline (requires GPU for stages 1 & 4)
    python zkaedi_audit_pipeline.py contract.sol --full

    # Stages 2+3 only (CPU, uses mock signatures)
    python zkaedi_audit_pipeline.py --preset defi_lending_pool

    # Custom signatures JSON → swarm + leviathan
    python zkaedi_audit_pipeline.py --signatures vulns.json

REQUIREMENTS:
    pip install gradio_client huggingface_hub safetensors numpy
    # For stages 1 & 4: pip install transformers torch accelerate
"""
from __future__ import annotations

import json
import time
import sys
import os
import argparse
import numpy as np
from pathlib import Path
from dataclasses import dataclass


# ══════════════════════════════════════════════════════════════
# ANSI COLORS
# ══════════════════════════════════════════════════════════════
class C:
    RST = "\033[0m"; B = "\033[1m"; D = "\033[2m"
    CY = "\033[96m"; MG = "\033[95m"; GR = "\033[92m"
    RD = "\033[91m"; YL = "\033[93m"; WH = "\033[97m"
    NC = "\033[38;2;0;255;255m"; NM = "\033[38;2;255;0;255m"
    NG = "\033[38;2;0;255;128m"; NK = "\033[38;2;255;215;0m"

    @staticmethod
    def grad(text, s, e):
        n = max(len(text), 1)
        return "".join(
            f"\033[38;2;{int(s[0]+(e[0]-s[0])*i/n)};{int(s[1]+(e[1]-s[1])*i/n)};{int(s[2]+(e[2]-s[2])*i/n)}m{c}"
            for i, c in enumerate(text)
        ) + C.RST


# ══════════════════════════════════════════════════════════════
# PIPELINE STAGES
# ══════════════════════════════════════════════════════════════
def stage_1_gemma_signatures(solidity_code: str, hf_token: str | None = None) -> list[dict]:
    """
    Stage 1: gemma-2-9b-solidity-merged
    Generates vulnerability energy signatures from Solidity source code.
    Requires GPU (or HF Inference Endpoint).
    """
    print(f"\n  {C.B}{C.NC}STAGE 1: VULNERABILITY SIGNATURE GENERATION{C.RST}")
    print(f"  {C.D}Model: zkaedi/gemma-2-9b-solidity-merged (9.2B params){C.RST}")

    try:
        from huggingface_hub import InferenceClient

        client = InferenceClient(
            model="zkaedi/gemma-2-9b-solidity-merged",
            token=hf_token,
        )

        prompt = f"""<start_of_turn>user
You are ZKAEDI PRIME — a Solidity smart contract energy signature auditor.
Analyze this contract and output vulnerability signatures as a JSON array.
Each signature must have: id, vuln_type, swc, severity, position [x,y],
energy, radius, description, compound_partner (if any), compound_type.

Contract:
```solidity
{solidity_code[:4000]}
```

Output ONLY a JSON array of vulnerability signatures, no other text.
<end_of_turn>
<start_of_turn>model
"""
        response = client.text_generation(
            prompt,
            max_new_tokens=2000,
            temperature=0.1,
            return_full_text=False,
        )

        # Parse JSON from response
        text = response.strip()
        if text.startswith("```"):
            text = text.split("```")[1]
            if text.startswith("json"):
                text = text[4:]
        signatures = json.loads(text)
        print(f"  {C.NG}Generated {len(signatures)} vulnerability signatures{C.RST}")
        return signatures

    except ImportError:
        print(f"  {C.YL}huggingface_hub InferenceClient not available{C.RST}")
        print(f"  {C.D}Falling back to preset signatures{C.RST}")
        return []
    except Exception as e:
        print(f"  {C.RD}Stage 1 failed: {e}{C.RST}")
        print(f"  {C.D}Falling back to preset signatures{C.RST}")
        return []


def stage_2_swarm_analysis(signatures: list[dict] | str,
                           steps: int = 300, window: int = 150,
                           seed: int = 42, preset: str | None = None) -> dict:
    """
    Stage 2: prime-swarm-hunter (Gradio Space)
    Runs 12-agent temporal compound detection.
    CPU only — calls the HF Space API.
    """
    print(f"\n  {C.B}{C.NC}STAGE 2: 12-AGENT SWARM COMPOUND DETECTION{C.RST}")
    print(f"  {C.D}Space: zkaedi/prime-swarm-hunter (12 agents, {steps} steps){C.RST}")

    try:
        from gradio_client import Client
        has_gradio = True
    except ImportError:
        has_gradio = False

    if has_gradio:
        try:
            client = Client("zkaedi/prime-swarm-hunter", verbose=False)

            if preset:
                print(f"  {C.D}Using preset: {preset}{C.RST}")
                result_json = client.predict(
                    preset, steps, window, seed,
                    api_name="/api_preset"
                )
            else:
                if isinstance(signatures, list):
                    signatures = json.dumps(signatures)
                print(f"  {C.D}Using {len(json.loads(signatures))} custom signatures{C.RST}")
                result_json = client.predict(
                    signatures, steps, window, seed,
                    api_name="/api_custom"
                )

            result = json.loads(result_json)
            s = result.get("summary", {})
            print(f"  {C.NG}Solo: {s.get('solo_detected', 0)}/{s.get('total_signatures', 0)} "
                  f"| Compound: {s.get('compounds_detected', 0)}/{s.get('compound_patterns', 0)} "
                  f"| Risk: {s.get('risk_score', 0)}{C.RST}")
            return result
        except Exception as e:
            print(f"  {C.YL}Space unavailable: {e}{C.RST}")
            has_gradio = False  # Fall through to local

    # Local fallback
    try:
        sys.path.insert(0, str(Path(__file__).parent))
        from prime_swarm_engine import run_preset, run_custom
        print(f"  {C.YL}Running swarm locally...{C.RST}")
        if preset:
            result = run_preset(preset, steps, window, seed)
        else:
            sigs = json.loads(signatures) if isinstance(signatures, str) else signatures
            result = run_custom(sigs, steps, window, seed)
        s = result.get("summary", {})
        print(f"  {C.NG}Solo: {s.get('solo_detected', 0)}/{s.get('total_signatures', 0)} "
              f"| Compound: {s.get('compounds_detected', 0)}/{s.get('compound_patterns', 0)} "
              f"| Risk: {s.get('risk_score', 0)}{C.RST}")
        return result
    except Exception as e2:
        print(f"  {C.RD}Local engine also failed: {e2}{C.RST}")
        return {"error": str(e2)}


def stage_3_leviathan_classify(swarm_result: dict, seed: int = 42) -> dict:
    """
    Stage 3: leviathan-v2
    CNN exploit topology classification + PRIME bistable refinement.
    Generates a synthetic manifold from swarm findings and classifies it.
    """
    print(f"\n  {C.B}{C.NC}STAGE 3: LEVIATHAN TOPOLOGY CLASSIFICATION{C.RST}")
    print(f"  {C.D}Model: zkaedi/leviathan-v2 (264K params, PRIME refinement){C.RST}")

    try:
        from huggingface_hub import hf_hub_download
        sys.path.insert(0, str(Path(__file__).parent))

        # Download and load Leviathan
        weights_path = hf_hub_download("zkaedi/leviathan-v2",
                                        "leviathan_v2_session_trained.safetensors")

        from safetensors.numpy import load_file
        weights = load_file(weights_path)

        # Import Leviathan class (try local, then download)
        try:
            from leviathan import Leviathan
        except ImportError:
            leviathan_py = hf_hub_download("zkaedi/leviathan-v2", "leviathan.py")
            import importlib.util
            spec = importlib.util.spec_from_file_location("leviathan", leviathan_py)
            mod = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(mod)
            Leviathan = mod.Leviathan

        model = Leviathan(weights)

        # Build synthetic manifold from swarm findings
        # Each detected vulnerability creates an energy signature in the manifold
        rng = np.random.default_rng(seed)
        H = rng.normal(0, 0.1, (256, 256)).astype(np.float32)
        V = np.zeros((256, 256), dtype=np.float32)

        solo_findings = swarm_result.get("solo_findings", [])
        compound_findings = swarm_result.get("compound_findings", [])

        # Inject vulnerability energy signatures into manifold
        for finding in solo_findings:
            pos = finding.get("position", [128, 128])
            energy = {"CRITICAL": 8.0, "HIGH": 5.0, "MEDIUM": 3.0, "LOW": 1.0}.get(
                finding.get("severity", "MEDIUM"), 3.0)
            x, y = int(pos[0] * 2.56), int(pos[1] * 2.56)  # Scale 0-100 → 0-256
            x, y = min(max(x, 5), 250), min(max(y, 5), 250)
            # Gaussian energy injection
            for dx in range(-10, 11):
                for dy in range(-10, 11):
                    nx, ny = x + dx, y + dy
                    if 0 <= nx < 256 and 0 <= ny < 256:
                        d = np.sqrt(dx**2 + dy**2)
                        H[nx, ny] += energy * np.exp(-d**2 / 50)

        # Compound findings amplify the manifold
        for comp in compound_findings:
            H *= 1.3  # Compound presence amplifies overall threat energy

        # Normalize
        if H.max() > 0:
            H = H / H.max()

        # Run Leviathan audit
        result = model.audit(H, V, seed=seed)
        vc = C.RD if result["verdict"] == "THREAT" else (
            C.NG if result["verdict"] == "CLEAN" else C.YL)
        print(f"  {vc}{C.B}Verdict: {result['verdict']}{C.RST}  "
              f"raw={result['raw_score']:.4f}  refined={result['refined_score']:.4f}  "
              f"committed={result['committed']}")

        return result

    except Exception as e:
        print(f"  {C.RD}Stage 3 failed: {e}{C.RST}")
        # Derive verdict from swarm risk score alone
        risk = swarm_result.get("summary", {}).get("risk_score", 0)
        verdict = "THREAT" if risk > 30 else ("UNCERTAIN" if risk > 10 else "CLEAN")
        return {"raw_score": risk / 100.0, "refined_score": risk / 100.0,
                "verdict": verdict, "committed": verdict != "UNCERTAIN",
                "prime_H": 0.0, "iterations": 0, "fallback": True}


def stage_4_final_report(solidity_code: str, swarm_result: dict,
                          leviathan_result: dict, hf_token: str | None = None) -> str:
    """
    Stage 4: solidity-vuln-auditor-7b
    Synthesizes a professional audit report from all findings.
    Requires GPU (or HF Inference Endpoint).
    """
    print(f"\n  {C.B}{C.NC}STAGE 4: PROFESSIONAL AUDIT REPORT{C.RST}")
    print(f"  {C.D}Model: zkaedi/solidity-vuln-auditor-7b (7.6B params, Qwen2){C.RST}")

    # Build report context from stages 2+3
    s = swarm_result.get("summary", {})
    findings_text = ""
    for f in swarm_result.get("solo_findings", []):
        findings_text += (f"- [{f.get('severity', 'MEDIUM')}] {f.get('swc', '')} "
                          f"{f.get('vuln_type', '')}: {f.get('description', '')}\n")
    for c in swarm_result.get("compound_findings", []):
        findings_text += (f"- [CRITICAL COMPOUND] {c.get('compound_type', '')}: "
                          f"{' + '.join(c.get('components', []))}\n")

    leviathan_verdict = leviathan_result.get("verdict", "UNCERTAIN")

    try:
        from huggingface_hub import InferenceClient

        client = InferenceClient(
            model="zkaedi/solidity-vuln-auditor-7b",
            token=hf_token,
        )

        prompt = f"""<|im_start|>user
Generate a professional smart contract security audit report.

SWARM ANALYSIS (12-agent Hamiltonian compound detection):
- Solo vulnerabilities: {s.get('solo_detected', 0)}/{s.get('total_signatures', 0)}
- Compound patterns: {s.get('compounds_detected', 0)}/{s.get('compound_patterns', 0)}
- Risk score: {s.get('risk_score', 0)}

FINDINGS:
{findings_text}

LEVIATHAN TOPOLOGY CLASSIFICATION:
- Verdict: {leviathan_verdict}
- Confidence: {leviathan_result.get('refined_score', 0):.4f}
- PRIME committed: {leviathan_result.get('committed', False)}

CONTRACT (first 2000 chars):
```solidity
{solidity_code[:2000]}
```

Write a professional audit report with executive summary, findings table,
risk assessment, and remediation recommendations.
<|im_end|>
<|im_start|>assistant
"""
        response = client.text_generation(
            prompt, max_new_tokens=2000, temperature=0.2, return_full_text=False)
        print(f"  {C.NG}Report generated ({len(response)} chars){C.RST}")
        return response

    except Exception as e:
        print(f"  {C.YL}Stage 4 model unavailable: {e}{C.RST}")
        print(f"  {C.D}Generating report from pipeline data...{C.RST}")

        # Fallback: generate report from pipeline data
        report = _generate_fallback_report(swarm_result, leviathan_result)
        return report


def _generate_fallback_report(swarm_result: dict, leviathan_result: dict) -> str:
    """Generate a structured report without the LLM."""
    s = swarm_result.get("summary", {})
    lines = [
        "=" * 60,
        "  ZKAEDI PRIME SECURITY AUDIT REPORT",
        "  Generated by: Unified Audit Pipeline v1.0",
        "=" * 60,
        "",
        "EXECUTIVE SUMMARY",
        "-" * 40,
        f"  Overall Verdict:    {leviathan_result.get('verdict', 'UNKNOWN')}",
        f"  Risk Score:         {s.get('risk_score', 0)}",
        f"  Solo Findings:      {s.get('solo_detected', 0)}/{s.get('total_signatures', 0)}",
        f"  Compound Patterns:  {s.get('compounds_detected', 0)}/{s.get('compound_patterns', 0)}",
        f"  Leviathan Score:    {leviathan_result.get('refined_score', 0):.4f}",
        f"  PRIME Committed:    {leviathan_result.get('committed', False)}",
        "",
        "FINDINGS",
        "-" * 40,
    ]

    for f in swarm_result.get("solo_findings", []):
        lines.append(f"  [{f.get('severity', '?'):<8}] {f.get('swc', 'N/A'):<8} "
                      f"{f.get('vuln_type', 'unknown')}")
        lines.append(f"             {f.get('description', '')}")
        lines.append(f"             Detected at step {f.get('detected_at_step', '?')} "
                      f"by {f.get('detected_by_role', '?')}")
        lines.append("")

    if swarm_result.get("compound_findings"):
        lines.append("COMPOUND VULNERABILITIES (CRITICAL)")
        lines.append("-" * 40)
        for c in swarm_result["compound_findings"]:
            lines.append(f"  {c.get('compound_type', 'unknown')}")
            lines.append(f"    Components: {' + '.join(c.get('components', []))}")
            lines.append(f"    Temporal gap: {c.get('temporal_gap', '?')} steps")
            agents = c.get("agents_involved", {})
            lines.append(f"    Discovered by: {agents.get('agent_a', {}).get('role', '?')} "
                          f"+ {agents.get('agent_b', {}).get('role', '?')}")
            lines.append("")

    lines.extend([
        "METHODOLOGY",
        "-" * 40,
        "  Stage 1: Vulnerability energy signature generation (Gemma 2 9B)",
        "  Stage 2: 12-agent Hamiltonian swarm with temporal correlation",
        "  Stage 3: Leviathan CNN topology classification + PRIME refinement",
        "  Stage 4: Report synthesis",
        "",
        "PIPELINE PARAMETERS",
        "-" * 40,
        f"  Swarm agents:       12",
        f"  Swarm steps:        {swarm_result.get('config', {}).get('steps', 300)}",
        f"  Temporal window:    {swarm_result.get('config', {}).get('temporal_window', 150)}",
        f"  PRIME eta:          3.50",
        f"  PRIME gamma:        0.30",
        f"  Leviathan params:   264,897",
        "",
        "=" * 60,
        "  ZKAEDI PRIME — This is not standard auditing.",
        "  This is computational physics meeting security analysis.",
        "=" * 60,
    ])

    return "\n".join(lines)


# ══════════════════════════════════════════════════════════════
# MAIN ORCHESTRATOR
# ══════════════════════════════════════════════════════════════
def run_pipeline(solidity_code: str | None = None,
                 signatures: list[dict] | None = None,
                 preset: str | None = None,
                 steps: int = 300,
                 window: int = 150,
                 seed: int = 42,
                 hf_token: str | None = None,
                 skip_stage_1: bool = False,
                 skip_stage_4: bool = False) -> dict:
    """
    Run the full ZKAEDI PRIME audit pipeline.

    Args:
        solidity_code: Raw Solidity source code (for stages 1 & 4)
        signatures: Pre-computed vulnerability signatures (skips stage 1)
        preset: Use a preset scenario instead of real code
        steps: Swarm simulation steps
        window: Temporal correlation window
        seed: Random seed
        hf_token: HuggingFace API token
        skip_stage_1: Skip Gemma signature generation
        skip_stage_4: Skip Qwen report generation
    """
    t_start = time.time()

    b = C.grad("=" * 65, (0, 255, 255), (255, 0, 255))
    print(f"\n{b}")
    print(C.grad("  🔱  ZKAEDI PRIME — UNIFIED SECURITY AUDIT PIPELINE  🔱  ",
                 (255, 0, 255), (0, 255, 255)))
    print(f"  {C.D}4-stage: Gemma → Swarm → Leviathan → Auditor{C.RST}")
    print(f"{b}\n")

    # ── Stage 1: Signature Generation ─────────────────────────
    if signatures:
        print(f"  {C.D}Stage 1: Skipped (signatures provided){C.RST}")
        sigs = signatures
    elif preset:
        print(f"  {C.D}Stage 1: Skipped (using preset: {preset}){C.RST}")
        sigs = None
    elif solidity_code and not skip_stage_1:
        sigs = stage_1_gemma_signatures(solidity_code, hf_token)
        if not sigs:
            print(f"  {C.YL}No signatures generated, falling back to preset{C.RST}")
            preset = "defi_lending_pool"
            sigs = None
    else:
        print(f"  {C.D}Stage 1: Skipped{C.RST}")
        if not preset:
            preset = "defi_lending_pool"
        sigs = None

    # ── Stage 2: Swarm Analysis ───────────────────────────────
    if preset:
        swarm_result = stage_2_swarm_analysis(None, steps, window, seed, preset=preset)
    else:
        swarm_result = stage_2_swarm_analysis(sigs, steps, window, seed)

    if "error" in swarm_result:
        print(f"  {C.RD}Pipeline aborted at Stage 2{C.RST}")
        return {"error": swarm_result["error"]}

    # ── Stage 3: Leviathan Classification ─────────────────────
    leviathan_result = stage_3_leviathan_classify(swarm_result, seed)

    # ── Stage 4: Report Generation ────────────────────────────
    if solidity_code and not skip_stage_4:
        report = stage_4_final_report(solidity_code, swarm_result, leviathan_result, hf_token)
    else:
        report = _generate_fallback_report(swarm_result, leviathan_result)

    elapsed = time.time() - t_start

    # ── Final Summary ─────────────────────────────────────────
    s = swarm_result.get("summary", {})
    verdict = leviathan_result.get("verdict", "UNKNOWN")
    vc = C.RD if verdict == "THREAT" else (C.NG if verdict == "CLEAN" else C.YL)

    print(f"\n{C.grad('=' * 65, (0, 255, 255), (255, 0, 255))}")
    print(f"  {C.B}{C.NK}AUDIT COMPLETE{C.RST}  {elapsed:.1f}s")
    print(f"\n  {C.B}Verdict:    {vc}{verdict}{C.RST}")
    print(f"  {C.B}Risk Score: {s.get('risk_score', 0)}{C.RST}")
    print(f"  {C.B}Solo:       {s.get('solo_detected', 0)}/{s.get('total_signatures', 0)}{C.RST}")
    print(f"  {C.B}Compound:   {s.get('compounds_detected', 0)}/{s.get('compound_patterns', 0)}{C.RST}")
    print(f"  {C.B}Leviathan:  {leviathan_result.get('refined_score', 0):.4f} "
          f"(committed={leviathan_result.get('committed', False)}){C.RST}")
    print(f"{C.grad('=' * 65, (255, 0, 255), (0, 255, 255))}\n")

    return {
        "verdict": verdict,
        "risk_score": s.get("risk_score", 0),
        "elapsed": elapsed,
        "swarm": swarm_result,
        "leviathan": leviathan_result,
        "report": report,
    }


# ══════════════════════════════════════════════════════════════
# CLI
# ══════════════════════════════════════════════════════════════
def main():
    parser = argparse.ArgumentParser(
        description="🔱 ZKAEDI PRIME Unified Security Audit Pipeline")
    parser.add_argument("contract", nargs="?", help="Solidity file path")
    parser.add_argument("--preset", choices=["defi_lending_pool", "nft_marketplace", "token_bridge"],
                        help="Use preset vulnerability scenario")
    parser.add_argument("--signatures", help="JSON file with vulnerability signatures")
    parser.add_argument("--steps", type=int, default=300, help="Swarm steps (default: 300)")
    parser.add_argument("--window", type=int, default=150, help="Temporal window (default: 150)")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument("--token", help="HuggingFace API token")
    parser.add_argument("--full", action="store_true", help="Run all 4 stages (requires GPU)")
    parser.add_argument("--report-only", action="store_true", help="Print report to stdout")
    parser.add_argument("--json", action="store_true", help="Output results as JSON")

    args = parser.parse_args()

    hf_token = args.token or os.environ.get("HF_TOKEN")
    solidity_code = None
    signatures = None

    if args.contract:
        with open(args.contract) as f:
            solidity_code = f.read()

    if args.signatures:
        with open(args.signatures) as f:
            signatures = json.load(f)

    result = run_pipeline(
        solidity_code=solidity_code,
        signatures=signatures,
        preset=args.preset or (None if (args.contract or args.signatures) else "defi_lending_pool"),
        steps=args.steps,
        window=args.window,
        seed=args.seed,
        hf_token=hf_token,
        skip_stage_1=not args.full,
        skip_stage_4=not args.full,
    )

    if args.json:
        # Clean for JSON serialization
        output = {k: v for k, v in result.items() if k != "report"}
        output["report_length"] = len(result.get("report", ""))
        print(json.dumps(output, indent=2, default=str))
    elif args.report_only:
        print(result.get("report", "No report generated"))
    else:
        print(result.get("report", ""))


if __name__ == "__main__":
    main()