package ai.intelliswarm.swarmai.skill.runtime;

import ai.intelliswarm.swarmai.skill.ContainerSkillSpec;
import ai.intelliswarm.swarmai.tool.base.BaseTool;

import java.time.Duration;
import java.util.List;
import java.util.Map;

/**
 * Factories for container-backed pentest tools. Each method returns a
 * {@link BaseTool} the agent can call directly — no orchestration code,
 * no shell-chain plans, no monolithic Kali images. The agent sees these
 * in its tool list and chooses which one to invoke based on what it
 * just learned.
 *
 * <p>Design principles:
 * <ul>
 *   <li><b>One tool, one job</b> — nmap_recon does discovery + service
 *       version, whatweb does HTTP fingerprinting, smbmap does SMB shares.
 *       Composing them is the agent's job, not the framework's.</li>
 *   <li><b>Focused images</b> — each tool ships in a 30–60 MB alpine
 *       image with just the binary it needs. No 900 MB Kali bloat. Tools
 *       that need a heavier base are explicit about it.</li>
 *   <li><b>Structured output</b> — every tool returns a parsed,
 *       agent-friendly summary. No raw stdout dumps that overflow the
 *       context window.</li>
 *   <li><b>Triggerwhen / avoidwhen</b> — every tool tells the agent when
 *       to pick it and when not to. The agent's tool-selection prompt
 *       ingests these directly.</li>
 * </ul>
 *
 * <h2>How the agent sees them</h2>
 * Each tool exposes:
 * <ul>
 *   <li>{@code getFunctionName()} — the tool name the agent calls</li>
 *   <li>{@code getDescription()} — what the tool does + expected output</li>
 *   <li>{@code getParameterSchema()} — JSON schema for arguments</li>
 *   <li>{@code getTriggerWhen()} / {@code getAvoidWhen()} — selection hints</li>
 * </ul>
 *
 * <h2>Adding a new tool</h2>
 * Write a new static factory here. Build a small Dockerfile, wrap in
 * {@link ContainerSkillSpec}, give the result a name + description +
 * output parser, return a {@link ContainerBackedBaseTool}.
 */
public final class PentestTools {

    private PentestTools() {}

    // ─── 1. nmap_recon — network discovery + service version detection ───
    public static BaseTool nmapRecon(ContainerSkillRuntime runtime) {
        ContainerSkillSpec spec = ContainerSkillSpec.builderFromDockerfile(
                "FROM alpine:3.19\n"
              + "RUN apk add --no-cache nmap nmap-scripts\n"
              + "ENTRYPOINT [\"sh\", \"-c\", \"nmap -F -sV --top-ports 100 -T4 $TARGET\"]\n")
                .envFromParam("target", "TARGET")
                .network("bridge")
                .memory("512m")
                .timeout(Duration.ofMinutes(5))
                .build();

        return ContainerBackedBaseTool.builder()
                .name("nmap_recon")
                .description("Run nmap port scan + service-version detection on a CIDR or single host. "
                        + "Probes the top 100 TCP ports, identifies running services, and returns "
                        + "a structured list of (host, port, service, version) findings.")
                .triggerWhen("First step of any network audit; needs an IP address, hostname, or CIDR range.")
                .avoidWhen("Don't run on hosts you don't own. Don't run on /16 or larger — too slow.")
                .category("security")
                .tags(List.of("network", "discovery", "port-scan", "service-version"))
                .parameterSchema(stringTargetSchema(
                        "Target CIDR, hostname, or IP (e.g. 192.168.1.0/24, 10.0.0.5, example.com)"))
                .spec(spec)
                .runtime(runtime)
                .outputParser(stdout -> {
                    List<Finding> findings = NmapOutputParser.parse(stdout);
                    return formatFindings(findings, "nmap_recon");
                })
                .build();
    }

    // ─── 2. whatweb — HTTP server + technology fingerprint for one host ───
    public static BaseTool whatweb(ContainerSkillRuntime runtime) {
        ContainerSkillSpec spec = ContainerSkillSpec.builderFromDockerfile(
                "FROM alpine:3.19\n"
              + "RUN apk add --no-cache whatweb curl\n"
              + "ENTRYPOINT [\"sh\", \"-c\", \"whatweb --quiet --no-errors -a 1 $TARGET || curl -sS -I $TARGET\"]\n")
                .envFromParam("target", "TARGET")
                .network("bridge")
                .memory("256m")
                .timeout(Duration.ofMinutes(2))
                .build();

        return ContainerBackedBaseTool.builder()
                .name("whatweb")
                .description("Fingerprint a web server: identifies HTTP server software, frameworks, "
                        + "CMS platforms, JavaScript libraries. Falls back to a curl HEAD when the "
                        + "site doesn't expose a banner.")
                .triggerWhen("After nmap_recon finds a host with port 80/443/8080 open and you want "
                        + "to identify what the web server is running.")
                .avoidWhen("Don't run against hosts that don't have a web server (waste of time). "
                        + "Don't run against the same target multiple times.")
                .category("security")
                .tags(List.of("web", "http", "fingerprint"))
                .parameterSchema(stringTargetSchema(
                        "Target URL or host (e.g. http://192.168.1.1, https://example.com)"))
                .spec(spec)
                .runtime(runtime)
                .outputParser(stdout -> {
                    String trimmed = stdout == null ? "" : stdout.trim();
                    if (trimmed.isEmpty()) return "(no fingerprint produced)";
                    if (trimmed.length() > 1500) trimmed = trimmed.substring(0, 1497) + "...";
                    return "whatweb fingerprint:\n" + trimmed;
                })
                .build();
    }

    // ─── 3. smbmap — list SMB shares + permissions on a host ───
    public static BaseTool smbmap(ContainerSkillRuntime runtime) {
        // smbmap isn't in the alpine community repo; pip-install in a small
        // python:3.12-alpine image keeps the footprint reasonable (~60MB).
        ContainerSkillSpec spec = ContainerSkillSpec.builderFromDockerfile(
                "FROM python:3.12-alpine\n"
              + "RUN apk add --no-cache gcc musl-dev libffi-dev && pip install --no-cache-dir smbmap\n"
              + "ENTRYPOINT [\"sh\", \"-c\", \"smbmap -H $TARGET -u guest -p '' 2>&1 || smbmap -H $TARGET\"]\n")
                .envFromParam("target", "TARGET")
                .network("bridge")
                .memory("256m")
                .timeout(Duration.ofMinutes(2))
                .build();

        return ContainerBackedBaseTool.builder()
                .name("smbmap")
                .description("Enumerate SMB shares + permissions on a target host. Tries guest "
                        + "credentials first, then anonymous. Returns the share name, access level, "
                        + "and any visible files.")
                .triggerWhen("After nmap_recon finds port 445 (microsoft-ds) open on a host.")
                .avoidWhen("Don't run on hosts that don't have SMB. Don't attempt brute-force "
                        + "credentials — guest/anonymous only.")
                .category("security")
                .tags(List.of("smb", "share-enumeration", "windows"))
                .parameterSchema(stringTargetSchema(
                        "Target host (e.g. 192.168.1.1) — must have SMB / port 445 open"))
                .spec(spec)
                .runtime(runtime)
                .outputParser(PentestTools::truncateForAgent)
                .build();
    }

    // ─── 4. ssh_audit — SSH algorithm / cipher / KEX inspection ───
    public static BaseTool sshAudit(ContainerSkillRuntime runtime) {
        ContainerSkillSpec spec = ContainerSkillSpec.builderFromDockerfile(
                "FROM python:3.12-alpine\n"
              + "RUN pip install --no-cache-dir ssh-audit\n"
              + "ENTRYPOINT [\"sh\", \"-c\", \"ssh-audit $TARGET\"]\n")
                .envFromParam("target", "TARGET")
                .network("bridge")
                .memory("128m")
                .timeout(Duration.ofMinutes(2))
                .build();

        return ContainerBackedBaseTool.builder()
                .name("ssh_audit")
                .description("Audit an SSH server's algorithms, ciphers, KEX, and host keys. "
                        + "Flags weak crypto and known-vulnerable configurations.")
                .triggerWhen("After nmap_recon finds port 22 (or any ssh service) open on a host.")
                .avoidWhen("Don't run on non-SSH hosts. The audit is read-only — no auth attempts.")
                .category("security")
                .tags(List.of("ssh", "crypto-audit"))
                .parameterSchema(stringTargetSchema(
                        "Target host (e.g. 192.168.1.5) — must have SSH / port 22 open"))
                .spec(spec)
                .runtime(runtime)
                .outputParser(PentestTools::truncateForAgent)
                .build();
    }

    // ─── 5. tls_check — certificate + cipher inspection for HTTPS endpoints ───
    public static BaseTool tlsCheck(ContainerSkillRuntime runtime) {
        ContainerSkillSpec spec = ContainerSkillSpec.builderFromDockerfile(
                "FROM alpine:3.19\n"
              + "RUN apk add --no-cache openssl\n"
              + "ENTRYPOINT [\"sh\", \"-c\", \"echo | openssl s_client -connect $TARGET -servername $TARGET 2>/dev/null | openssl x509 -noout -subject -issuer -dates -ext subjectAltName 2>/dev/null\"]\n")
                .envFromParam("target", "TARGET")
                .network("bridge")
                .memory("64m")
                .timeout(Duration.ofMinutes(1))
                .build();

        return ContainerBackedBaseTool.builder()
                .name("tls_check")
                .description("Inspect the TLS certificate of an HTTPS endpoint: subject, issuer, "
                        + "validity dates, and Subject Alternative Names. Read-only — no probes "
                        + "for weak ciphers (ssh_audit covers SSH; this is the TLS analogue).")
                .triggerWhen("After nmap_recon finds port 443 (or any ssl/tls service) open.")
                .avoidWhen("Don't run on plain HTTP hosts. Provide host:port form (e.g. 192.168.1.1:443).")
                .category("security")
                .tags(List.of("tls", "certificate", "https"))
                .parameterSchema(stringTargetSchema(
                        "host:port form (e.g. 192.168.1.1:443) — must have TLS open"))
                .spec(spec)
                .runtime(runtime)
                .outputParser(PentestTools::truncateForAgent)
                .build();
    }

    // ─── 6. dns_recon — basic DNS enumeration via dig ───
    public static BaseTool dnsRecon(ContainerSkillRuntime runtime) {
        ContainerSkillSpec spec = ContainerSkillSpec.builderFromDockerfile(
                "FROM alpine:3.19\n"
              + "RUN apk add --no-cache bind-tools\n"
              + "ENTRYPOINT [\"sh\", \"-c\", \"dig +short ANY $TARGET; dig +short NS $TARGET; dig +short MX $TARGET\"]\n")
                .envFromParam("target", "TARGET")
                .network("bridge")
                .memory("64m")
                .timeout(Duration.ofMinutes(1))
                .build();

        return ContainerBackedBaseTool.builder()
                .name("dns_recon")
                .description("Run DNS lookups on a domain: ANY records, NS, MX. Useful for "
                        + "understanding a domain's infrastructure (mail servers, name servers).")
                .triggerWhen("When the target is a domain name and you want DNS-level metadata. "
                        + "NOT useful against raw IPs or CIDRs (dig won't return anything).")
                .avoidWhen("Don't run against IP addresses or CIDR ranges — dig expects a domain.")
                .category("security")
                .tags(List.of("dns", "domain-recon"))
                .parameterSchema(stringTargetSchema(
                        "Domain name (e.g. example.com) — NOT an IP address"))
                .spec(spec)
                .runtime(runtime)
                .outputParser(PentestTools::truncateForAgent)
                .build();
    }

    /** Convenience: assemble the full library in one call. */
    public static List<BaseTool> all(ContainerSkillRuntime runtime) {
        return List.of(
                nmapRecon(runtime),
                whatweb(runtime),
                smbmap(runtime),
                sshAudit(runtime),
                tlsCheck(runtime),
                dnsRecon(runtime));
    }

    // ---------------------------------------------------------------------
    // Helpers
    // ---------------------------------------------------------------------

    /** A simple "string target" parameter schema reused by every tool. */
    private static Map<String, Object> stringTargetSchema(String description) {
        return Map.of(
                "type", "object",
                "properties", Map.of("target", Map.of(
                        "type", "string",
                        "description", description)),
                "required", List.of("target"));
    }

    /**
     * Format a finding list as a compact text block the agent can read.
     * Keeps the structure but doesn't bury the agent in JSON noise.
     */
    private static String formatFindings(List<Finding> findings, String toolName) {
        if (findings == null || findings.isEmpty()) {
            return toolName + ": no findings";
        }
        StringBuilder sb = new StringBuilder(toolName)
                .append(": ").append(findings.size()).append(" finding(s)\n");
        for (Finding f : findings) {
            sb.append("  - ").append(f.summary()).append("\n");
        }
        return sb.toString();
    }

    /** Trim raw tool output to a size that fits an agent's context window. */
    private static String truncateForAgent(String stdout) {
        if (stdout == null) return "(no output)";
        String s = stdout.trim();
        if (s.isEmpty()) return "(no output)";
        if (s.length() > 1500) s = s.substring(0, 1497) + "...";
        return s;
    }
}
