"""Smart-default policy generator.

Takes scan results and generates a starter agentward.yaml policy with
security-aware defaults. Users review and customize the generated policy.

The generator applies per-tool rules based on risk analysis:
  - CRITICAL/shell tools → require_approval
  - Write-capable tools → restrict to read-only where safe
  - Cross-server risk combinations → skill_chaining rules
  - Network+credentials → outbound blocking
"""

from __future__ import annotations

from pathlib import Path

import yaml

from agentward.policy.schema import (
    AgentWardPolicy,
    ChainingRule,
    ResourcePermissions,
)
from agentward.scan.permissions import (
    DataAccessType,
    RiskLevel,
    ScanResult,
    ServerPermissionMap,
    ToolPermission,
)

_POLICY_VERSION = "1.0"

# Header comment prepended to generated YAML
_YAML_HEADER = """\
# AgentWard Policy — auto-generated by `agentward configure`
#
# This policy was generated from a scan of your tools and skills.
# Review each section and adjust to your needs:
#
#   skills:            Per-tool data access permissions
#   skill_chaining:    Cross-tool execution restrictions
#   require_approval:  Tools that need human confirmation before running
#
# Apply this policy: agentward setup --policy agentward.yaml
# Re-generate:       agentward configure [target]
# Documentation:     https://agentward.ai/docs/policy-format
#
"""


def generate_policy(scan: ScanResult) -> AgentWardPolicy:
    """Generate a smart-default policy from scan results.

    Applies security-aware rules to each tool based on its risk analysis:
      - CRITICAL/shell/destructive tools → require_approval
      - Write-capable email/messaging → restrict to read-only
      - Write-capable filesystem → restrict to read-only
      - Network+credentials on same server → block outbound
      - Cross-server chaining risks → skill_chaining rules

    Args:
        scan: Complete scan result from build_permission_map().

    Returns:
        A valid AgentWardPolicy ready to serialize to YAML.
    """
    skills: dict[str, dict[str, ResourcePermissions]] = {}
    require_approval: list[str] = []
    chaining_rules: list[ChainingRule] = []

    for server_map in scan.servers:
        server_name = server_map.server.name
        server_resources: dict[str, ResourcePermissions] = {}

        for tool_perm in server_map.tools:
            _apply_tool_rules(
                tool_perm, server_name, server_resources, require_approval,
            )

        # Server-level rules: network+credentials → block outbound
        _apply_server_rules(server_map, server_resources)

        if server_resources:
            skills[server_name] = server_resources

    # Cross-server chaining rules
    chaining_rules = _compute_chaining_rules(scan.servers)

    # Deduplicate require_approval
    seen: set[str] = set()
    deduped_approval: list[str] = []
    for tool_name in require_approval:
        if tool_name not in seen:
            deduped_approval.append(tool_name)
            seen.add(tool_name)

    return AgentWardPolicy(
        version=_POLICY_VERSION,
        skills=skills,
        skill_chaining=chaining_rules,
        require_approval=deduped_approval,
    )


def serialize_policy(policy: AgentWardPolicy) -> str:
    """Serialize a policy to YAML with header comments.

    Args:
        policy: The policy to serialize.

    Returns:
        A YAML string with explanatory header comments.
    """
    # Build a clean dict for YAML output
    data = _policy_to_dict(policy)

    yaml_body = yaml.dump(
        data,
        default_flow_style=False,
        sort_keys=False,
        allow_unicode=True,
    )

    return _YAML_HEADER + yaml_body


def write_policy(policy: AgentWardPolicy, path: Path) -> None:
    """Write a policy to a YAML file.

    Args:
        policy: The policy to write.
        path: Output file path.

    Raises:
        PermissionError: If the path is not writable.
    """
    content = serialize_policy(policy)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content, encoding="utf-8")


# ---------------------------------------------------------------------------
# Per-tool rules
# ---------------------------------------------------------------------------


def _apply_tool_rules(
    tool_perm: ToolPermission,
    server_name: str,
    resources: dict[str, ResourcePermissions],
    require_approval: list[str],
) -> None:
    """Apply security rules for a single tool.

    Args:
        tool_perm: The tool's permission analysis.
        server_name: Which server this tool belongs to.
        resources: Dict to populate with resource restrictions (mutated).
        require_approval: List to append approval-gated tools to (mutated).
    """
    access_types = {a.type for a in tool_perm.data_access}
    tool_name = tool_perm.tool.name

    # Rule 1: CRITICAL risk or shell access → require approval
    if tool_perm.risk_level == RiskLevel.CRITICAL or DataAccessType.SHELL in access_types:
        require_approval.append(tool_name)

    # Rule 2: Destructive tools → require approval
    if tool_perm.is_destructive:
        require_approval.append(tool_name)

    # Rule 3: Write-capable email/messaging → restrict to read-only
    if not tool_perm.is_read_only:
        if DataAccessType.EMAIL in access_types:
            resource_key = _infer_resource_key(tool_name, "email")
            resources[resource_key] = _make_resource(
                read=True, send=False, delete=False, draft=True,
            )

        if DataAccessType.MESSAGING in access_types:
            resource_key = _infer_resource_key(tool_name, "messaging")
            resources[resource_key] = _make_resource(
                read=True, send=False, delete=False,
            )

    # Rule 4: Write-capable filesystem → restrict to read-only
    if not tool_perm.is_read_only and DataAccessType.FILESYSTEM in access_types:
        resource_key = _infer_resource_key(tool_name, "filesystem")
        resources[resource_key] = _make_resource(read=True, write=False)


def _apply_server_rules(
    server_map: ServerPermissionMap,
    resources: dict[str, ResourcePermissions],
) -> None:
    """Apply server-level rules based on cross-tool analysis.

    Args:
        server_map: The server's permission map.
        resources: Dict to populate with resource restrictions (mutated).
    """
    server_access: set[DataAccessType] = set()
    for tool_perm in server_map.tools:
        for access in tool_perm.data_access:
            server_access.add(access.type)

    # Network + credentials on same server → deny each network-capable tool.
    # We emit restrictions keyed by resource names derived from the actual
    # tool names so the policy engine can match them (instead of a generic
    # "network" key that no tool name would ever match).
    if DataAccessType.NETWORK in server_access and DataAccessType.CREDENTIALS in server_access:
        for tool_perm in server_map.tools:
            tool_types = {a.type for a in tool_perm.data_access}
            if DataAccessType.NETWORK in tool_types:
                resource_key = _infer_resource_key(tool_perm.tool.name, tool_perm.tool.name)
                resources[resource_key] = _make_resource(outbound=False)


# ---------------------------------------------------------------------------
# Cross-server chaining rules
# ---------------------------------------------------------------------------


def _compute_chaining_rules(
    servers: list[ServerPermissionMap],
) -> list[ChainingRule]:
    """Detect cross-server chaining risks and generate rules.

    Checks for dangerous combinations:
      - Email + browser (prompt injection via URLs)
      - Data source + shell (prompt injection → code execution)

    Args:
        servers: All server permission maps.

    Returns:
        A list of chaining rules to include in the policy.
    """
    rules: list[ChainingRule] = []

    # Collect per-server capabilities
    server_caps: list[tuple[str, set[DataAccessType], bool]] = []
    for server_map in servers:
        types: set[DataAccessType] = set()
        has_shell = False
        for tool in server_map.tools:
            for access in tool.data_access:
                types.add(access.type)
            if DataAccessType.SHELL in {a.type for a in tool.data_access}:
                has_shell = True
        server_caps.append((server_map.server.name, types, has_shell))

    # Check pairs
    seen_rules: set[tuple[str, str]] = set()

    for i, (name_a, types_a, shell_a) in enumerate(server_caps):
        for name_b, types_b, shell_b in server_caps[i + 1:]:
            # Email + browser → block chaining
            if DataAccessType.EMAIL in types_a and DataAccessType.BROWSER in types_b:
                key = (name_a, name_b)
                if key not in seen_rules:
                    rules.append(ChainingRule(
                        source_skill=name_a, target_skill=name_b,
                    ))
                    seen_rules.add(key)
            if DataAccessType.EMAIL in types_b and DataAccessType.BROWSER in types_a:
                key = (name_b, name_a)
                if key not in seen_rules:
                    rules.append(ChainingRule(
                        source_skill=name_b, target_skill=name_a,
                    ))
                    seen_rules.add(key)

            # Data source + shell → block chaining
            if shell_b and not shell_a:
                key = (name_a, name_b)
                if key not in seen_rules:
                    rules.append(ChainingRule(
                        source_skill=name_a, target_skill=name_b,
                    ))
                    seen_rules.add(key)
            if shell_a and not shell_b:
                key = (name_b, name_a)
                if key not in seen_rules:
                    rules.append(ChainingRule(
                        source_skill=name_b, target_skill=name_a,
                    ))
                    seen_rules.add(key)

    return rules


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_resource(**actions: bool) -> ResourcePermissions:
    """Create a ResourcePermissions bypassing the model_validator.

    The model_validator on ResourcePermissions expects flat YAML-style
    dicts (e.g., ``{"read": True, "send": False}``), not the structured
    ``actions=`` keyword. Constructing directly avoids the validator
    mis-interpreting the ``actions`` key.

    Args:
        **actions: Action name → allowed boolean pairs.

    Returns:
        A ResourcePermissions with the given action permissions.
    """
    obj = ResourcePermissions.model_construct(
        denied=False,
        actions=dict(actions),
        filters={},
    )
    return obj


def _infer_resource_key(tool_name: str, default: str) -> str:
    """Infer a resource key from a tool name.

    Tries to extract a meaningful resource name from the tool name
    (e.g., "gmail_send" → "gmail", "slack_post" → "slack").
    Falls back to the provided default.

    Args:
        tool_name: The tool name.
        default: Fallback resource key.

    Returns:
        A resource key string for the policy YAML.
    """
    from agentward.scan.permissions import (
        _READ_VERBS,
        _WRITE_VERBS,
        _DELETE_VERBS,
        _EXECUTE_VERBS,
    )

    all_verbs = _READ_VERBS | _WRITE_VERBS | _DELETE_VERBS | _EXECUTE_VERBS

    for sep in ("_", "-", "."):
        if sep in tool_name:
            parts = tool_name.split(sep)
            for part in parts:
                if part.lower() not in all_verbs and len(part) > 1:
                    return part.lower()

    return default


def _policy_to_dict(policy: AgentWardPolicy) -> dict:
    """Convert a policy to a plain dict suitable for yaml.dump().

    Produces cleaner YAML than pydantic's model_dump() by:
    - Omitting empty sections
    - Converting ChainingRule to readable strings
    - Flattening ResourcePermissions to action dicts

    Args:
        policy: The policy to convert.

    Returns:
        A dict ready for yaml.dump().
    """
    data: dict = {"version": policy.version}

    # Skills section
    if policy.skills:
        skills_dict: dict = {}
        for skill_name, resources in policy.skills.items():
            resource_dict: dict = {}
            for resource_name, perms in resources.items():
                if perms.denied:
                    resource_dict[resource_name] = {"denied": True}
                elif perms.actions:
                    resource_dict[resource_name] = dict(perms.actions)
                    if perms.filters:
                        resource_dict[resource_name]["filters"] = dict(perms.filters)
            if resource_dict:
                skills_dict[skill_name] = resource_dict
        if skills_dict:
            data["skills"] = skills_dict

    # Require approval
    if policy.require_approval:
        data["require_approval"] = list(policy.require_approval)

    # Skill chain depth
    if policy.skill_chain_depth is not None:
        data["skill_chain_depth"] = policy.skill_chain_depth

    # Chaining mode (include if chaining rules exist)
    if policy.skill_chaining:
        data["chaining_mode"] = policy.chaining_mode.value

    # Skill chaining
    if policy.skill_chaining:
        data["skill_chaining"] = [
            f"{rule.source_skill} cannot trigger {rule.target_skill}"
            for rule in policy.skill_chaining
        ]

    # Data boundaries
    if policy.data_boundaries:
        boundaries: dict = {}
        for name, boundary in policy.data_boundaries.items():
            boundaries[name] = {
                "skills": list(boundary.skills),
                "classification": boundary.classification,
                "rules": list(boundary.rules),
                "on_violation": boundary.on_violation.value,
            }
        data["data_boundaries"] = boundaries

    return data
