"""
ZIA scaffold templates used by `zia init`.

These are the template strings for YAML + Python agent class boilerplate per domain.
"""

from __future__ import annotations

from typing import NamedTuple


class DomainTemplate(NamedTuple):
    yaml_template: str
    python_template: str
    default_tools: list[str]
    default_allowed_actions: list[str]
    default_denied_actions: list[str]


DOMAIN_TEMPLATES: dict[str, DomainTemplate] = {
    "risk_quant": DomainTemplate(
        yaml_template="""\
agent:
  id: {agent_id}
  name: "{agent_name}"
  domain: risk_quant
  version: "1.0.0"

intent:
  goal: "Compute and update risk scores for critical assets"
  success_criteria:
    - "All assets have updated risk scores in the SIF graph"
  priority: high

reasoning:
  mode: probabilistic
  autonomy_level: bounded
  confidence_threshold: 0.75

capabilities:
  tools:
    - read_asset
    - list_assets
    - list_vulnerabilities
    - compute_risk
    - write_risk_node
  data_access:
    - asset_inventory
    - vulnerability_scanner

boundaries:
  risk_budget: medium
  allowed_actions:
    - agent_execute
    - read_asset
    - list_assets
    - compute_risk
    - write_risk_node
  denied_actions:
    - delete_asset
    - delete_vulnerability
  environment_scope:
    - production
    - staging

safety:
  guardrails:
    - no_destructive_actions
    - require_confidence_threshold
    - redact_secrets
  sandbox_profile: standard
  audit_level: standard
""",
        python_template="""\
\"\"\"
{agent_name} — ZIA Risk Quantification Agent.

Auto-generated by `zia init`. Implement execute() to define your agent's logic.
Register with: @register_agent(domain="risk_quant")
\"\"\"

from __future__ import annotations

from zia.core.runtime.agent import AgentContext, AgentResult, BaseAgent
from zia.core.runtime.registry import register_agent
from zia.core.tools.substrate import ToolExecutor
import zia.core.tools.builtins as tools  # noqa: F401


@register_agent(
    domain="risk_quant",
    description="{agent_name}",
    version="1.0.0",
)
class {class_name}(BaseAgent):
    \"\"\"
    {agent_name}

    Uses the ZIA risk propagation engine to compute and write risk scores.
    \"\"\"

    def execute(self, context: AgentContext) -> AgentResult:
        # 1. Load assets
        assets = ToolExecutor.call(tools.list_assets, context=context)

        results = []
        for asset in assets:
            # 2. Compute risk
            risk = ToolExecutor.call(
                tools.compute_risk,
                context=context,
                criticality=asset.get("criticality", "medium"),
                exposure=asset.get("exposure_level", "internal"),
            )
            results.append({{
                "asset_id": asset["node_id"],
                **risk,
            }})

        return AgentResult.ok(context, output={{"results": results}})
""",
        default_tools=["read_asset", "list_assets", "list_vulnerabilities", "compute_risk", "write_risk_node"],
        default_allowed_actions=["read_asset", "list_assets", "compute_risk", "write_risk_node"],
        default_denied_actions=["delete_asset"],
    ),

    "supply_chain": DomainTemplate(
        yaml_template="""\
agent:
  id: {agent_id}
  name: "{agent_name}"
  domain: supply_chain
  version: "1.0.0"

intent:
  goal: "Assess and update risk scores for third-party vendors"
  success_criteria:
    - "All tier-1 vendors have updated risk scores"
  priority: high

reasoning:
  mode: hybrid
  autonomy_level: bounded
  confidence_threshold: 0.70

capabilities:
  tools:
    - list_vendors
    - write_risk_node
  data_access:
    - vendor_registry

boundaries:
  risk_budget: medium
  allowed_actions:
    - list_vendors
    - write_vendor_risk_score
    - write_risk_node
  denied_actions:
    - delete_vendor
    - modify_contract
  environment_scope:
    - production
    - staging
  approval_gates:
    - flag_critical_vendor

safety:
  guardrails:
    - no_destructive_actions
    - redact_secrets
  sandbox_profile: standard
  audit_level: standard
""",
        python_template="""\
\"\"\"
{agent_name} — ZIA Supply Chain Risk Agent.

Auto-generated by `zia init`. Implement execute() to define your agent's logic.
\"\"\"

from __future__ import annotations

from zia.core.runtime.agent import AgentContext, AgentResult, BaseAgent
from zia.core.runtime.registry import register_agent
from zia.core.tools.substrate import ToolExecutor
import zia.core.tools.builtins as tools  # noqa: F401


@register_agent(
    domain="supply_chain",
    description="{agent_name}",
    version="1.0.0",
)
class {class_name}(BaseAgent):
    \"\"\"
    {agent_name}

    Assesses third-party vendor risk and writes findings to the SIF graph.
    \"\"\"

    def execute(self, context: AgentContext) -> AgentResult:
        vendors = ToolExecutor.call(tools.list_vendors, context=context)

        results = []
        for vendor in vendors:
            # TODO: implement your vendor scoring logic here
            score = self._score_vendor(vendor)
            results.append({{"vendor_id": vendor["node_id"], "score": score}})

        return AgentResult.ok(context, output={{"vendors_assessed": len(results), "results": results}})

    def _score_vendor(self, vendor: dict) -> float:
        \"\"\"Implement your vendor risk scoring logic here.\"\"\"
        tier = int(vendor.get("tier", 2))
        return {{1: 6.0, 2: 4.0, 3: 2.0}}.get(tier, 4.0)
""",
        default_tools=["list_vendors", "write_risk_node"],
        default_allowed_actions=["list_vendors", "write_vendor_risk_score", "write_risk_node"],
        default_denied_actions=["delete_vendor"],
    ),

    "red_team": DomainTemplate(
        yaml_template="""\
agent:
  id: {agent_id}
  name: "{agent_name}"
  domain: red_team
  version: "1.0.0"

intent:
  goal: "Discover and document attack paths in the target environment"
  success_criteria:
    - "All critical attack paths identified"
    - "Blast radius estimated for each finding"
  priority: critical

reasoning:
  mode: hybrid
  autonomy_level: bounded
  confidence_threshold: 0.85

capabilities:
  tools:
    - scan_network
    - probe_service
    - capture_evidence
  data_access:
    - asset_inventory

boundaries:
  risk_budget: high
  allowed_actions:
    - scan_network
    - probe_service
    - capture_evidence
  denied_actions:
    - delete_data
    - exfiltrate_data
    - persist_access
  environment_scope:
    - staging
  approval_gates:
    - execute_exploit
    - escalate_privileges

safety:
  guardrails:
    - blast_radius_limit
    - no_data_exfiltration
    - environment_scope_enforcement
    - redact_secrets
  sandbox_profile: offensive_isolated
  audit_level: verbose
""",
        python_template="""\
\"\"\"
{agent_name} — ZIA Red Team Agent.

Auto-generated by `zia init`. Implement execute() to define your agent's logic.
IMPORTANT: Red team agents must only target environments in environment_scope (default: staging).
\"\"\"

from __future__ import annotations

from zia.core.runtime.agent import AgentContext, AgentResult, BaseAgent
from zia.core.runtime.registry import register_agent


@register_agent(
    domain="red_team",
    description="{agent_name}",
    version="1.0.0",
)
class {class_name}(BaseAgent):
    \"\"\"
    {agent_name}

    Discovers attack paths within bounded, sandboxed environments.
    \"\"\"

    def execute(self, context: AgentContext) -> AgentResult:
        # Implement red team logic here
        # Always check context.environment before any targeted action
        findings: list[dict] = []

        # TODO: implement your attack path discovery logic

        return AgentResult.ok(context, output={{"findings": findings}})
""",
        default_tools=["scan_network", "probe_service", "capture_evidence"],
        default_allowed_actions=["scan_network", "probe_service", "capture_evidence"],
        default_denied_actions=["delete_data", "exfiltrate_data"],
    ),

    "appsec": DomainTemplate(
        yaml_template="""\
agent:
  id: {agent_id}
  name: "{agent_name}"
  domain: appsec
  version: "1.0.0"

intent:
  goal: "Scan code and dependencies for security vulnerabilities"
  success_criteria:
    - "All critical/high findings identified and written to SIF"
  priority: high

reasoning:
  mode: rule_based
  autonomy_level: bounded
  confidence_threshold: 0.80

capabilities:
  tools:
    - sast_scan
    - dependency_scan
    - write_vulnerability_node
  data_access:
    - source_code
    - dependency_manifest

boundaries:
  risk_budget: low
  allowed_actions:
    - sast_scan
    - dependency_scan
    - write_vulnerability_node
  denied_actions:
    - modify_source_code
    - delete_dependency
  environment_scope:
    - staging
    - dev

safety:
  guardrails:
    - no_destructive_actions
    - redact_secrets
  sandbox_profile: standard
  audit_level: standard
""",
        python_template="""\
\"\"\"
{agent_name} — ZIA AppSec Agent.

Auto-generated by `zia init`. Implement execute() to define your agent's logic.
\"\"\"

from __future__ import annotations

from zia.core.runtime.agent import AgentContext, AgentResult, BaseAgent
from zia.core.runtime.registry import register_agent


@register_agent(
    domain="appsec",
    description="{agent_name}",
    version="1.0.0",
)
class {class_name}(BaseAgent):
    \"\"\"
    {agent_name}

    Performs SAST/DAST scanning and writes findings as VulnerabilityNodes to the SIF graph.
    \"\"\"

    def execute(self, context: AgentContext) -> AgentResult:
        findings: list[dict] = []

        # TODO: implement your SAST/DAST/dependency scanning logic here
        # Use ToolExecutor.call(tools.write_vulnerability_node, ...) to persist findings

        return AgentResult.ok(context, output={{"findings_count": len(findings), "findings": findings}})
""",
        default_tools=["sast_scan", "dependency_scan", "write_vulnerability_node"],
        default_allowed_actions=["sast_scan", "dependency_scan", "write_vulnerability_node"],
        default_denied_actions=["modify_source_code"],
    ),

    "ai_security": DomainTemplate(
        yaml_template="""\
agent:
  id: {agent_id}
  name: "{agent_name}"
  domain: ai_security
  version: "1.0.0"

intent:
  goal: "Assess AI/LLM model risk including prompt injection and data leakage"
  success_criteria:
    - "All AI models have updated risk scores"
    - "Prompt injection risks documented"
  priority: high

reasoning:
  mode: hybrid
  autonomy_level: bounded
  confidence_threshold: 0.80

capabilities:
  tools:
    - prompt_risk_scan
    - model_exposure_check
    - write_ai_model_node
  data_access:
    - ai_model_registry

boundaries:
  risk_budget: medium
  allowed_actions:
    - prompt_risk_scan
    - model_exposure_check
    - write_ai_model_node
  denied_actions:
    - modify_model
    - delete_model
  environment_scope:
    - production
    - staging

safety:
  guardrails:
    - no_destructive_actions
    - prompt_injection_defense
    - redact_secrets
  sandbox_profile: standard
  audit_level: verbose
""",
        python_template="""\
\"\"\"
{agent_name} — ZIA AI/LLM Security Agent.

Auto-generated by `zia init`. Implement execute() to define your agent's logic.
\"\"\"

from __future__ import annotations

from zia.core.runtime.agent import AgentContext, AgentResult, BaseAgent
from zia.core.runtime.registry import register_agent


@register_agent(
    domain="ai_security",
    description="{agent_name}",
    version="1.0.0",
)
class {class_name}(BaseAgent):
    \"\"\"
    {agent_name}

    Assesses prompt injection risks, model exposure, and data leakage for AI systems.
    \"\"\"

    def execute(self, context: AgentContext) -> AgentResult:
        findings: list[dict] = []

        # TODO: implement prompt risk analysis, model exposure detection here

        return AgentResult.ok(context, output={{"findings": findings}})
""",
        default_tools=["prompt_risk_scan", "model_exposure_check", "write_ai_model_node"],
        default_allowed_actions=["prompt_risk_scan", "model_exposure_check", "write_ai_model_node"],
        default_denied_actions=["modify_model"],
    ),

    "compliance": DomainTemplate(
        yaml_template="""\
agent:
  id: {agent_id}
  name: "{agent_name}"
  domain: compliance
  version: "1.0.0"

intent:
  goal: "Monitor controls, collect evidence, and detect compliance drift"
  success_criteria:
    - "All controls assessed against framework requirements"
    - "Evidence collected and stored"
  priority: medium

reasoning:
  mode: rule_based
  autonomy_level: bounded
  confidence_threshold: 0.85

capabilities:
  tools:
    - check_control
    - collect_evidence
    - write_control_node
  data_access:
    - control_framework
    - audit_logs

boundaries:
  risk_budget: low
  allowed_actions:
    - check_control
    - collect_evidence
    - write_control_node
  denied_actions:
    - delete_evidence
    - modify_control
  environment_scope:
    - production
    - staging

safety:
  guardrails:
    - no_destructive_actions
    - redact_secrets
  sandbox_profile: strict
  audit_level: verbose
""",
        python_template="""\
\"\"\"
{agent_name} — ZIA Continuous Compliance Agent.

Auto-generated by `zia init`. Implement execute() to define your agent's logic.
\"\"\"

from __future__ import annotations

from zia.core.runtime.agent import AgentContext, AgentResult, BaseAgent
from zia.core.runtime.registry import register_agent


@register_agent(
    domain="compliance",
    description="{agent_name}",
    version="1.0.0",
)
class {class_name}(BaseAgent):
    \"\"\"
    {agent_name}

    Monitors security controls and collects compliance evidence against a framework.
    \"\"\"

    def execute(self, context: AgentContext) -> AgentResult:
        control_results: list[dict] = []

        # TODO: implement control monitoring and evidence collection here

        return AgentResult.ok(
            context,
            output={{
                "controls_checked": len(control_results),
                "results": control_results,
            }},
        )
""",
        default_tools=["check_control", "collect_evidence", "write_control_node"],
        default_allowed_actions=["check_control", "collect_evidence", "write_control_node"],
        default_denied_actions=["delete_evidence", "modify_control"],
    ),
}

SUPPORTED_DOMAINS = list(DOMAIN_TEMPLATES.keys())
