from __future__ import annotations

from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Tuple

import yaml

from bastion.models import ScanResult

POLICY_FILE = "policy.yaml"


def load_policy(bastion_dir: Path) -> Dict[str, Any]:
    """Load policy.yaml from .bastion directory if it exists, otherwise return empty structure."""
    policy_path = bastion_dir / POLICY_FILE
    if not policy_path.exists():
        return {"acknowledged": {}}
    with open(policy_path, "r") as f:
        data = yaml.safe_load(f)
    return data if data is not None else {"acknowledged": {}}


def save_policy(policy: Dict[str, Any], bastion_dir: Path) -> None:
    """Save the policy dict to .bastion/policy.yaml."""
    bastion_dir.mkdir(parents=True, exist_ok=True)
    policy_path = bastion_dir / POLICY_FILE
    with open(policy_path, "w") as f:
        yaml.dump(policy, f, default_flow_style=False, sort_keys=False)


def add_acknowledgment(bastion_dir: Path, plugin_id: str, category: str, reason: str) -> None:
    """Add an acknowledgment to the policy, avoiding duplicates."""
    policy = load_policy(bastion_dir)
    
    if "acknowledged" not in policy:
        policy["acknowledged"] = {}
    
    if plugin_id not in policy["acknowledged"]:
        policy["acknowledged"][plugin_id] = []
    
    # Check if this plugin_id + category already exists
    for ack in policy["acknowledged"][plugin_id]:
        if ack.get("category") == category:
            return  # Already exists, don't add duplicate
    
    # Add new acknowledgment
    acknowledged_at = datetime.now(timezone.utc).isoformat()
    policy["acknowledged"][plugin_id].append({
        "category": category,
        "reason": reason,
        "acknowledged_at": acknowledged_at,
    })
    
    save_policy(policy, bastion_dir)


def remove_acknowledgment(bastion_dir: Path, plugin_id: str, category: str) -> bool:
    """Remove an acknowledgment, returns True if found and removed."""
    policy = load_policy(bastion_dir)
    
    if plugin_id not in policy.get("acknowledged", {}):
        return False
    
    original_length = len(policy["acknowledged"][plugin_id])
    policy["acknowledged"][plugin_id] = [
        ack for ack in policy["acknowledged"][plugin_id]
        if ack.get("category") != category
    ]
    
    # If list is now empty, remove the plugin_id key
    if not policy["acknowledged"][plugin_id]:
        del policy["acknowledged"][plugin_id]
    
    found_and_removed = len(policy["acknowledged"][plugin_id]) < original_length if plugin_id in policy["acknowledged"] else original_length > 0
    
    if found_and_removed:
        save_policy(policy, bastion_dir)
    
    return found_and_removed


def is_acknowledged(policy: Dict[str, Any], plugin_id: str, category: str) -> bool:
    """Returns True if the given plugin_id + category is acknowledged in the policy."""
    if plugin_id not in policy.get("acknowledged", {}):
        return False
    
    for ack in policy["acknowledged"][plugin_id]:
        if ack.get("category") == category:
            return True
    
    return False


def filter_acknowledged(scan_results: List[ScanResult], policy: Dict[str, Any]) -> Tuple[List[ScanResult], int]:
    """
    Filter scan results based on the policy.
    
    Returns a tuple of (unacknowledged_results, acknowledged_count).
    - unacknowledged_results: Original list of ScanResult objects unchanged
    - acknowledged_count: Count of capabilities across all results that are acknowledged
    """
    acknowledged_count = 0
    
    for result in scan_results:
        for capability in result.capabilities:
            if is_acknowledged(policy, result.plugin_id, capability.category):
                acknowledged_count += 1
    
    return (scan_results, acknowledged_count)
