"""
Gemini function-calling declarations for the ReportAgent tools.
"""
from google.genai import types as genai_types  # type: ignore[attr-defined]

TOOL_DECLARATIONS = [
    genai_types.FunctionDeclaration(
        name="tavily_search",
        description=(
            "Search the web for current information about cyber "
            "threats, vulnerabilities, or security topics."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search query string.",
                },
                "max_results": {
                    "type": "integer",
                    "description": (
                        "Maximum number of results to return "
                        "(default 5)."
                    ),
                },
            },
            "required": ["query"],
        },
    ),
    genai_types.FunctionDeclaration(
        name="nvd_search",
        description=(
            "Search the NIST National Vulnerability Database for "
            "CVE entries by CVE ID or keyword."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "cve_id": {
                    "type": "string",
                    "description": (
                        "A specific CVE identifier "
                        "(e.g. CVE-2021-44228)."
                    ),
                },
                "keyword": {
                    "type": "string",
                    "description": (
                        "A keyword to search for in CVE "
                        "descriptions."
                    ),
                },
            },
        },
    ),
    genai_types.FunctionDeclaration(
        name="mitre_search",
        description=(
            "Search the MITRE ATT&CK framework for techniques "
            "by technique ID or keyword."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "technique_id": {
                    "type": "string",
                    "description": (
                        "An ATT&CK technique ID (e.g. T1059)."
                    ),
                },
                "search": {
                    "type": "string",
                    "description": (
                        "A keyword to search in technique "
                        "names and descriptions."
                    ),
                },
            },
        },
    ),
    genai_types.FunctionDeclaration(
        name="virustotal_scan",
        description=(
            "Look up an indicator on VirusTotal (IP address, "
            "domain, URL, or file hash)."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "scan_type": {
                    "type": "string",
                    "description": (
                        "The type of indicator: ip, domain, "
                        "url, or hash."
                    ),
                },
                "value": {
                    "type": "string",
                    "description": (
                        "The indicator value to look up."
                    ),
                },
            },
            "required": ["scan_type", "value"],
        },
    ),
    genai_types.FunctionDeclaration(
        name="abuseipdb_check",
        description=(
            "Check an IP address against the AbuseIPDB database "
            "for abuse reports and reputation."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "ip": {
                    "type": "string",
                    "description": "The IP address to check.",
                },
            },
            "required": ["ip"],
        },
    ),
    genai_types.FunctionDeclaration(
        name="otx_search",
        description=(
            "Search AlienVault OTX for threat intelligence on "
            "an indicator (IP, domain, hash, CVE, etc.)."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "indicator_type": {
                    "type": "string",
                    "description": (
                        "The indicator type: IPv4, IPv6, "
                        "domain, hostname, url, hash, or cve."
                    ),
                },
                "value": {
                    "type": "string",
                    "description": (
                        "The indicator value to look up."
                    ),
                },
            },
            "required": ["indicator_type", "value"],
        },
    ),
    genai_types.FunctionDeclaration(
        name="dt_exec",
        description=(
            "Execute a shell command on a digital-twin "
            "container. Use this to inspect processes, "
            "network connections, logs, services, and "
            "file systems on DT hosts. Valid container "
            "names: i1_gateway, i1_firewall, i1_log_collector, "
            "i1_server_1–i1_server_6 (Incident 1) or "
            "i2_server_1–i2_server_6 (Incident 2)."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "container": {
                    "type": "string",
                    "description": (
                        "The host id of the container "
                        "(e.g. i1_firewall, i1_server_1)."
                    ),
                },
                "command": {
                    "type": "string",
                    "description": (
                        "The shell command to execute."
                    ),
                },
            },
            "required": ["container", "command"],
        },
    ),
    genai_types.FunctionDeclaration(
        name="dt_python_exec",
        description=(
            "Execute a Python script in an isolated "
            "sandbox container. Use this for analysis "
            "scripts, data parsing, or computations "
            "that are easier to express in Python. "
            "The sandbox is NOT connected to the "
            "digital twin — you cannot call dt_exec, "
            "subprocess, or access DT containers "
            "from within this sandbox. Only use it "
            "to process/analyze data you have already "
            "collected via dt_exec."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": (
                        "The Python source code to run."
                    ),
                },
            },
            "required": ["code"],
        },
    ),
    genai_types.FunctionDeclaration(
        name="generate_attack_image",
        description=(
            "Generate a visual attack path diagram using AI "
            "image generation, overlaid on the network "
            "topology. Call this once you have gathered "
            "enough evidence to illustrate the attack path, "
            "but before producing the final assessment."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string",
                    "description": (
                        "A detailed description of the "
                        "attack path to illustrate, "
                        "including which hosts were "
                        "compromised, in what order, and "
                        "what protocols/techniques were "
                        "used at each step."
                    ),
                },
            },
            "required": ["prompt"],
        },
    ),
    genai_types.FunctionDeclaration(
        name="dt_restart",
        description=(
            "Restart a digital-twin container that has "
            "crashed or stopped. Use when dt_exec fails "
            "with a 'container is not running' error. "
            "Pass a specific container name to restart "
            "just that host, or pass 'all' to redeploy "
            "the entire digital twin."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "container": {
                    "type": "string",
                    "description": (
                        "The host id to restart "
                        "(e.g. i1_server_2), or 'all' "
                        "to redeploy the entire DT."
                    ),
                },
            },
            "required": ["container"],
        },
    ),
    genai_types.FunctionDeclaration(
        name="run_host_analyzers",
        description=(
            "Delegate deep analysis of one or more hosts "
            "to parallel HostAnalyzerAgent sub-agents. "
            "Each sub-agent will independently investigate "
            "the specified host using DT shell commands, "
            "external lookups, and other tools, then "
            "produce a structured host analysis report. "
            "Use this when multiple hosts need detailed "
            "investigation — only include hosts that are "
            "relevant to the incident."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "hosts": {
                    "type": "array",
                    "description": (
                        "List of hosts to analyze in "
                        "parallel."
                    ),
                    "items": {
                        "type": "object",
                        "properties": {
                            "host_id": {
                                "type": "string",
                                "description": (
                                    "The container / "
                                    "host identifier "
                                    "(e.g. i1_server_1)."
                                ),
                            },
                            "host_description": {
                                "type": "string",
                                "description": (
                                    "A brief description "
                                    "of why this host "
                                    "is relevant and "
                                    "what to look for."
                                ),
                            },
                        },
                        "required": [
                            "host_id",
                            "host_description",
                        ],
                    },
                },
            },
            "required": ["hosts"],
        },
    ),
    genai_types.FunctionDeclaration(
        name="produce_assessment",
        description=(
            "Produce the final structured incident assessment. "
            "Call this ONLY after gathering enough information "
            "from multiple investigation tools."
        ),
        parameters={  # type: ignore[arg-type]
            "type": "object",
            "properties": {
                "incident_summary": {
                    "type": "string",
                    "description": (
                        "Brief overview of the incident."
                    ),
                },
                "attack_vector_analysis": {
                    "type": "string",
                    "description": (
                        "How the attacker gained access and "
                        "what techniques were used."
                    ),
                },
                "indicators_of_compromise": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "type": {
                                "type": "string",
                                "description": (
                                    "IOC type: ip, domain, "
                                    "hash, cve, or other."
                                ),
                            },
                            "value": {
                                "type": "string",
                                "description": (
                                    "The IOC value."
                                ),
                            },
                            "context": {
                                "type": "string",
                                "description": (
                                    "Context about this IOC."
                                ),
                            },
                        },
                        "required": [
                            "type", "value", "context",
                        ],
                    },
                },
                "severity": {
                    "type": "string",
                    "description": (
                        "Severity: Critical, High, "
                        "Medium, or Low."
                    ),
                },
                "severity_justification": {
                    "type": "string",
                    "description": (
                        "Why this severity was chosen."
                    ),
                },
                "affected_assets": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "asset": {
                                "type": "string",
                                "description": (
                                    "The asset name."
                                ),
                            },
                            "impact": {
                                "type": "string",
                                "description": (
                                    "Impact on this asset."
                                ),
                            },
                        },
                        "required": ["asset", "impact"],
                    },
                },
            },
            "required": [
                "incident_summary",
                "attack_vector_analysis",
                "indicators_of_compromise",
                "severity",
                "severity_justification",
                "affected_assets",
            ],
        },
    ),
]
