"""
IntelliScan — Diagram Generator for Conception & Design Document
================================================================

Generates all architecture and design diagrams as PNG images.
Uses matplotlib + graphviz (via subprocess) + mermaid-py (optional).

Requirements:
    pip install matplotlib graphviz pillow

Usage:
    python docs/generate_diagrams.py

Output:
    docs/images/*.png
"""

import os
import subprocess
import textwrap
from pathlib import Path

# ── Output directory ─────────────────────────────────────────────────────
OUT_DIR = Path(__file__).resolve().parent / "images"
OUT_DIR.mkdir(parents=True, exist_ok=True)


# ═══════════════════════════════════════════════════════════════════════════
# DIAGRAM 1 — High-Level Pipeline (Graphviz DOT)
# ═══════════════════════════════════════════════════════════════════════════
def generate_pipeline_diagram():
    """6-module sequential pipeline."""
    dot = textwrap.dedent(r'''
    digraph Pipeline {
        rankdir=LR;
        bgcolor="#1a1a2e";
        node [shape=box, style="filled,rounded", fontname="Helvetica-Bold",
              fontsize=12, fontcolor="white", color="#16213e", penwidth=2];
        edge [color="#e94560", penwidth=2, fontname="Helvetica", fontsize=10,
              fontcolor="#a8a8b3"];

        Target [label="Target\nWeb App", fillcolor="#0f3460", shape=cylinder];

        M1 [label="Module 1\nCrawler", fillcolor="#1a1a6c"];
        M2 [label="Module 2\nInjector", fillcolor="#4a1a6c"];
        M3 [label="Module 3\nAnalyzer", fillcolor="#6c1a4a"];
        M4 [label="Module 4\nML Classifier", fillcolor="#6c3a1a"];
        M5 [label="Module 5\nPayload Gen", fillcolor="#1a6c3a"];
        M6 [label="Module 6\nReporter", fillcolor="#1a4a6c"];

        F1 [label="results.json", shape=note, fillcolor="#e94560",
            fontsize=9, fontcolor="white"];
        F2 [label="injection_\nresults.json", shape=note, fillcolor="#e94560",
            fontsize=9, fontcolor="white"];
        F3 [label="labeled_\nresults.json", shape=note, fillcolor="#e94560",
            fontsize=9, fontcolor="white"];
        F4 [label="model.pkl", shape=note, fillcolor="#e94560",
            fontsize=9, fontcolor="white"];
        F6 [label="report.pdf", shape=note, fillcolor="#e94560",
            fontsize=9, fontcolor="white"];

        Target -> M1 [label=" HTTP"];
        M1 -> F1;
        F1 -> M2;
        M2 -> F2;
        F2 -> M3;
        M3 -> F3;
        F3 -> M4;
        M4 -> F4;
        F3 -> M6;
        M6 -> F6;
        M5 -> M2 [label=" payloads", style=dashed, color="#53a8b6"];
    }
    ''')
    _render_dot(dot, "01_pipeline_diagram")


# ═══════════════════════════════════════════════════════════════════════════
# DIAGRAM 2 — Module Dependency Graph
# ═══════════════════════════════════════════════════════════════════════════
def generate_dependency_graph():
    dot = textwrap.dedent(r'''
    digraph Dependencies {
        rankdir=TB;
        bgcolor="#1a1a2e";
        node [shape=component, style="filled,rounded", fontname="Helvetica-Bold",
              fontsize=11, fontcolor="white", penwidth=2];
        edge [color="#53a8b6", penwidth=1.5, fontname="Helvetica", fontsize=9,
              fontcolor="#a8a8b3"];

        subgraph cluster_core {
            label="Core Layer"; fontname="Helvetica-Bold"; fontsize=13;
            fontcolor="#e94560"; color="#2a2a4e"; style="dashed";
            core [label="core.py\nIntelliScan", fillcolor="#0f3460"];
            config [label="config.py\nConstants", fillcolor="#16213e"];
        }

        subgraph cluster_modules {
            label="Modules Layer"; fontname="Helvetica-Bold"; fontsize=13;
            fontcolor="#e94560"; color="#2a2a4e"; style="dashed";
            crawler    [label="crawler.py",    fillcolor="#1a1a6c"];
            injector   [label="injector.py",   fillcolor="#4a1a6c"];
            analyzer   [label="analyzer.py",   fillcolor="#6c1a4a"];
            classifier [label="classifier.py", fillcolor="#6c3a1a"];
            payload    [label="payload_gen.py", fillcolor="#1a6c3a"];
            reporter   [label="reporter.py",   fillcolor="#1a4a6c"];
        }

        subgraph cluster_utils {
            label="Utils Layer"; fontname="Helvetica-Bold"; fontsize=13;
            fontcolor="#e94560"; color="#2a2a4e"; style="dashed";
            http     [label="http_client.py\nHttpClient", fillcolor="#3a3a5c"];
            notifier [label="notifier.py\nDiscordNotifier", fillcolor="#3a3a5c"];
        }

        subgraph cluster_cli {
            label="Interface Layer"; fontname="Helvetica-Bold"; fontsize=13;
            fontcolor="#e94560"; color="#2a2a4e"; style="dashed";
            cli [label="__main__.py\nCLI (Click)", fillcolor="#533a6c"];
        }

        cli -> core;
        core -> crawler; core -> injector; core -> analyzer;
        core -> classifier; core -> payload; core -> reporter;
        core -> notifier; core -> config;
        crawler -> http; crawler -> config;
        injector -> http; injector -> config;
        analyzer -> config;
        classifier -> config;
        payload -> config;
        reporter -> config;
        http -> config;
    }
    ''')
    _render_dot(dot, "02_dependency_graph")


# ═══════════════════════════════════════════════════════════════════════════
# DIAGRAM 3 — Data Flow Diagram
# ═══════════════════════════════════════════════════════════════════════════
def generate_data_flow():
    dot = textwrap.dedent(r'''
    digraph DataFlow {
        rankdir=TB;
        bgcolor="#1a1a2e";
        node [fontname="Helvetica", fontsize=11, fontcolor="white", penwidth=2];
        edge [color="#53a8b6", penwidth=1.5, fontname="Helvetica", fontsize=9,
              fontcolor="#a8a8b3"];

        // External entities
        user   [label="Security\nAnalyst", shape=doublecircle, fillcolor="#0f3460",
                style=filled];
        target [label="Target\nWeb App", shape=doublecircle, fillcolor="#0f3460",
                style=filled];
        discord[label="Discord\nWebhook", shape=doublecircle, fillcolor="#0f3460",
                style=filled];

        // Processes
        node [shape=ellipse, style=filled];
        P1 [label="1. Crawl\nBFS Exploration", fillcolor="#1a1a6c"];
        P2 [label="2. Inject\nPayload Delivery", fillcolor="#4a1a6c"];
        P3 [label="3. Analyze\nSignature Matching", fillcolor="#6c1a4a"];
        P4 [label="4. Classify\nRandom Forest ML", fillcolor="#6c3a1a"];
        P5 [label="5. Generate\nPayload Mutations", fillcolor="#1a6c3a"];
        P6 [label="6. Report\nPDF Generation", fillcolor="#1a4a6c"];

        // Data stores
        node [shape=cylinder, style=filled, fillcolor="#e94560"];
        D1 [label="results.json"];
        D2 [label="injection_results.json"];
        D3 [label="labeled_results.json"];
        D4 [label="model.pkl"];
        D5 [label="all_payloads.json"];
        D6 [label="report.pdf"];
        D7 [label="payloads/*.txt", fillcolor="#a84060"];

        // Flows
        user -> P1 [label="target URL\ncredentials"];
        P1 -> target [label="HTTP GET", style=dashed];
        target -> P1 [label="HTML pages", style=dashed];
        P1 -> D1 [label="forms, params"];

        D1 -> P2;
        D7 -> P2 [label="base payloads"];
        P2 -> target [label="GET/POST\ninjections", style=dashed];
        target -> P2 [label="responses", style=dashed];
        P2 -> D2;

        D2 -> P3;
        P3 -> D3 [label="labels +\nseverity"];

        D3 -> P4;
        P4 -> D4 [label="trained model"];

        D7 -> P5;
        P5 -> D5 [label="mutated\npayloads"];

        D3 -> P6;
        P6 -> D6;
        D6 -> user [label="PDF report"];
        P6 -> discord [label="notification", style=dashed];
    }
    ''')
    _render_dot(dot, "03_data_flow_diagram")


# ═══════════════════════════════════════════════════════════════════════════
# DIAGRAM 4 — ML Feature Importance (Matplotlib bar chart)
# ═══════════════════════════════════════════════════════════════════════════
def generate_ml_feature_chart():
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import matplotlib.ticker as mtick
    except ImportError:
        print("[SKIP] matplotlib not installed — skipping ML feature chart")
        return

    features = [
        ("response_len",              0.550),
        ("vuln_type_encoded",         0.335),
        ("payload_len",               0.116),
        ("response_to_payload_ratio", 0.030),
        ("response_length_anomaly",   0.025),
        ("payload_special_chars",     0.020),
        ("has_html_tags",             0.015),
        ("status_code",               0.009),
    ]
    names = [f[0] for f in features]
    values = [f[1] for f in features]
    colors = ["#e94560", "#e94560", "#e94560", "#53a8b6", "#53a8b6",
              "#53a8b6", "#53a8b6", "#53a8b6"]

    fig, ax = plt.subplots(figsize=(10, 5), facecolor="#1a1a2e")
    ax.set_facecolor("#1a1a2e")
    bars = ax.barh(names[::-1], values[::-1], color=colors[::-1], edgecolor="#2a2a4e",
                   height=0.6)
    ax.set_xlabel("Feature Importance", color="white", fontsize=12)
    ax.set_title("Random Forest — Feature Importances (No Data Leakage)",
                 color="white", fontsize=14, fontweight="bold", pad=15)
    ax.tick_params(colors="white", labelsize=10)
    ax.xaxis.set_major_formatter(mtick.PercentFormatter(1.0))
    for spine in ax.spines.values():
        spine.set_color("#2a2a4e")

    for bar, val in zip(bars, values[::-1]):
        ax.text(bar.get_width() + 0.005, bar.get_y() + bar.get_height() / 2,
                f"{val*100:.1f}%", va="center", color="white", fontsize=9)

    plt.tight_layout()
    path = OUT_DIR / "04_ml_feature_importance.png"
    fig.savefig(path, dpi=200, facecolor=fig.get_facecolor())
    plt.close(fig)
    print(f"[OK] {path}")


# ═══════════════════════════════════════════════════════════════════════════
# DIAGRAM 5 — Class / Component Diagram
# ═══════════════════════════════════════════════════════════════════════════
def generate_class_diagram():
    dot = textwrap.dedent(r'''
    digraph ClassDiagram {
        rankdir=TB;
        bgcolor="#1a1a2e";
        node [shape=record, style=filled, fontname="Courier", fontsize=10,
              fontcolor="white", fillcolor="#16213e", color="#53a8b6", penwidth=1.5];
        edge [color="#e94560", penwidth=1.5, fontname="Helvetica", fontsize=9,
              fontcolor="#a8a8b3"];

        Crawler [label="{Crawler|+ target: str\l+ auth: tuple\l+ max_pages: int\l+ max_depth: int\l+ http: HttpClient\l|+ run() : CrawlResult\l- _login()\l- _set_dvwa_security_low()\l- _bfs(start)\l- _extract_forms()\l- _extract_links()\l}"];

        CrawlResult [label="{CrawlResult|+ target: str\l+ pages_visited: int\l+ forms: list[Form]\l+ url_params: list[UrlParam]\l|+ to_dict() : dict\l+ save(path) : Path\l}"];

        Injector [label="{Injector|+ http: HttpClient\l+ target: str\l+ concurrent: bool\l+ results: list\l|+ run_all() : list[InjectionResult]\l- _build_tasks()\l- _inject_one()\l- _load_payloads()\l}"];

        Analyzer [label="{Analyzer|+ raw: list[dict]\l+ labeled: list[LabeledResult]\l|+ run() : list[LabeledResult]\l- _detect_sqli()\l- _detect_xss()\l- _detect_lfi()\l+ stats_by_type() : dict\l}"];

        MLClassifier [label="{MLClassifier|+ model: RandomForestClassifier\l+ is_trained: bool\l|+ train(dataset) : TrainingReport\l+ predict(entry) : tuple\l+ extract_features(entry) : list\l+ save(path)\l+ load(path)\l}"];

        PayloadGen [label="{PayloadGenerator|+ rng: Random\l|+ generate(payload, type) : list\l+ generate_all(payloads) : list\l- _case_mutation()\l- _comment_insertion()\l- _quote_swap()\l- _url_encode()\l- _double_url_encode()\l- _whitespace_variation()\l}"];

        Reporter [label="{Reporter (FPDF)|+ target: str\l+ results: list[dict]\l|+ build(output) : Path\l+ cover_page()\l+ section_summary()\l+ section_findings()\l+ section_recommendations()\l+ section_methodology()\l}"];

        HttpClient [label="{HttpClient|+ base_url: str\l+ session: Session\l+ timeout: int\l|+ get(url) : Response\l+ post(url) : Response\l+ get_csrf_token(url) : str\l+ close()\l}"];

        IntelliScan [label="{IntelliScan|+ target: str\l+ auth: tuple\l+ train_ml: bool\l|+ run() : ScanResult\l}"];

        IntelliScan -> Crawler [label="creates"];
        IntelliScan -> Injector [label="creates"];
        IntelliScan -> Analyzer [label="creates"];
        IntelliScan -> MLClassifier [label="creates"];
        IntelliScan -> PayloadGen [label="creates"];
        IntelliScan -> Reporter [label="creates"];
        Crawler -> HttpClient [label="uses"];
        Injector -> HttpClient [label="uses"];
        Crawler -> CrawlResult [label="returns"];
    }
    ''')
    _render_dot(dot, "05_class_diagram")


# ═══════════════════════════════════════════════════════════════════════════
# DIAGRAM 6 — Deployment / Docker Architecture
# ═══════════════════════════════════════════════════════════════════════════
def generate_deployment_diagram():
    dot = textwrap.dedent(r'''
    digraph Deployment {
        rankdir=TB;
        bgcolor="#1a1a2e";
        compound=true;
        node [fontname="Helvetica", fontsize=11, fontcolor="white", penwidth=2];
        edge [color="#53a8b6", penwidth=1.5, fontname="Helvetica", fontsize=9,
              fontcolor="#a8a8b3"];

        subgraph cluster_docker {
            label="Docker Compose — intelliscan-net (bridge)";
            fontname="Helvetica-Bold"; fontsize=13;
            fontcolor="#e94560"; color="#e94560"; style="dashed";
            bgcolor="#16213e";

            subgraph cluster_app {
                label="intelliscan container"; fontname="Helvetica-Bold";
                fontsize=11; fontcolor="#53a8b6"; color="#53a8b6";
                style="rounded,dashed"; bgcolor="#1a1a3e";

                flask [label="Flask\nDashboard\n:5000", shape=box,
                       style="filled,rounded", fillcolor="#1a6c3a"];
                cli_app [label="CLI\n(Click+Rich)", shape=box,
                       style="filled,rounded", fillcolor="#4a1a6c"];
                pipeline [label="6-Module\nPipeline", shape=box,
                       style="filled,rounded", fillcolor="#0f3460"];
            }

            subgraph cluster_dvwa {
                label="dvwa container"; fontname="Helvetica-Bold";
                fontsize=11; fontcolor="#53a8b6"; color="#53a8b6";
                style="rounded,dashed"; bgcolor="#1a1a3e";

                apache [label="Apache+PHP\n:80", shape=box,
                        style="filled,rounded", fillcolor="#6c1a4a"];
                mysql [label="MySQL", shape=cylinder,
                       style=filled, fillcolor="#6c3a1a"];
            }
        }

        // Volumes
        vol_results [label="./results\n(volume)", shape=folder,
                     style=filled, fillcolor="#e94560", fontcolor="white"];
        vol_models  [label="./models\n(volume)", shape=folder,
                     style=filled, fillcolor="#e94560", fontcolor="white"];

        // External
        user [label="Security\nAnalyst", shape=doublecircle,
              style=filled, fillcolor="#0f3460"];
        discord_ext [label="Discord\nWebhook", shape=doublecircle,
                     style=filled, fillcolor="#0f3460"];

        // Connections
        user -> flask [label=":5000"];
        user -> cli_app [label="CLI"];
        cli_app -> pipeline;
        flask -> pipeline;
        pipeline -> apache [label="HTTP\n:80"];
        apache -> mysql;
        pipeline -> vol_results;
        pipeline -> vol_models;
        pipeline -> discord_ext [label="webhook", style=dashed];
    }
    ''')
    _render_dot(dot, "06_deployment_diagram")


# ═══════════════════════════════════════════════════════════════════════════
# DIAGRAM 7 — ML Training / Evaluation Flow
# ═══════════════════════════════════════════════════════════════════════════
def generate_ml_flow():
    dot = textwrap.dedent(r'''
    digraph MLFlow {
        rankdir=LR;
        bgcolor="#1a1a2e";
        node [shape=box, style="filled,rounded", fontname="Helvetica",
              fontsize=11, fontcolor="white", penwidth=2];
        edge [color="#53a8b6", penwidth=1.5, fontname="Helvetica", fontsize=9,
              fontcolor="#a8a8b3"];

        D [label="labeled_results\n.json", shape=cylinder, fillcolor="#e94560"];
        FE [label="Feature\nExtraction\n(8 features)", fillcolor="#1a1a6c"];
        SPL [label="Train/Test\nSplit\n80/20\nStratified", fillcolor="#4a1a6c"];
        TR [label="Random Forest\nTraining\n(100 trees)", fillcolor="#6c3a1a"];
        EV [label="Evaluation\nAccuracy, F1\nConfusion Matrix", fillcolor="#1a6c3a"];
        CV [label="5-Fold\nCross-Val", fillcolor="#1a4a6c"];
        M [label="model.pkl", shape=cylinder, fillcolor="#e94560"];
        REP [label="TrainingReport\n(dataclass)", fillcolor="#6c1a4a"];

        D -> FE;
        FE -> SPL [label="X, y"];
        SPL -> TR [label="X_train\ny_train"];
        SPL -> EV [label="X_test\ny_test"];
        TR -> EV [label="y_pred"];
        FE -> CV [label="X, y"];
        TR -> M [label="joblib.dump"];
        EV -> REP;
        CV -> REP [label="cv_scores"];
    }
    ''')
    _render_dot(dot, "07_ml_training_flow")


# ═══════════════════════════════════════════════════════════════════════════
# DIAGRAM 8 — Payload Mutation Pipeline
# ═══════════════════════════════════════════════════════════════════════════
def generate_mutation_diagram():
    dot = textwrap.dedent(r'''
    digraph Mutations {
        rankdir=TB;
        bgcolor="#1a1a2e";
        node [shape=box, style="filled,rounded", fontname="Helvetica",
              fontsize=10, fontcolor="white", penwidth=1.5];
        edge [color="#53a8b6", penwidth=1.5];

        base [label="Base Payload\ne.g. ' OR 1=1 --", fillcolor="#e94560",
              fontsize=12, shape=ellipse];

        m1 [label="1. Case Alternation\noR 1=1 --", fillcolor="#1a1a6c"];
        m2 [label="2. SQL Comment\nOR/**/1=1 --", fillcolor="#4a1a6c"];
        m3 [label="3. Quote Swap\n\" OR 1=1 --", fillcolor="#6c1a4a"];
        m4 [label="4. URL Encode\n%27%20OR%201%3D1%20--", fillcolor="#6c3a1a"];
        m5 [label="5. Double Encode\n%2527%2520OR...", fillcolor="#1a6c3a"];
        m6 [label="6. Whitespace\n'  OR  1=1  --", fillcolor="#1a4a6c"];

        out [label="Deduplicated\nVariant Set\n~16x expansion", fillcolor="#e94560",
             fontsize=12, shape=ellipse];

        base -> m1; base -> m2; base -> m3;
        base -> m4; base -> m5; base -> m6;
        m1 -> out; m2 -> out; m3 -> out;
        m4 -> out; m5 -> out; m6 -> out;
    }
    ''')
    _render_dot(dot, "08_mutation_pipeline")


# ═══════════════════════════════════════════════════════════════════════════
# Helper — render DOT via graphviz CLI
# ═══════════════════════════════════════════════════════════════════════════
def _render_dot(dot_source: str, name: str):
    dot_file = OUT_DIR / f"{name}.dot"
    png_file = OUT_DIR / f"{name}.png"
    dot_file.write_text(dot_source, encoding="utf-8")
    try:
        subprocess.run(
            ["dot", "-Tpng", "-Gdpi=200", str(dot_file), "-o", str(png_file)],
            check=True, capture_output=True, text=True,
        )
        print(f"[OK] {png_file}")
    except FileNotFoundError:
        print(f"[WARN] Graphviz 'dot' not found. DOT source saved: {dot_file}")
        print("       Install: https://graphviz.org/download/")
        print(f"       Then run: dot -Tpng -Gdpi=200 {dot_file} -o {png_file}")
    except subprocess.CalledProcessError as e:
        print(f"[ERROR] dot failed for {name}: {e.stderr}")


# ═══════════════════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════════════════
if __name__ == "__main__":
    print("=" * 60)
    print("  IntelliScan — Diagram Generator")
    print("=" * 60)
    print(f"Output directory: {OUT_DIR}\n")

    generate_pipeline_diagram()
    generate_dependency_graph()
    generate_data_flow()
    generate_ml_feature_chart()
    generate_class_diagram()
    generate_deployment_diagram()
    generate_ml_flow()
    generate_mutation_diagram()

    print("\n" + "=" * 60)
    print("  Done! Check docs/images/ for outputs.")
    print("  DOT files saved for manual rendering if graphviz missing.")
    print("=" * 60)