"""MITRE ATT&CK mapper — maps extracted TTPs to ATT&CK technique IDs."""
import json
import hashlib
import logging
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

ATTACK_DATA_PATH = Path(__file__).parent.parent / 'data' / 'mitre_attack.json'


class MITREMapper:
    """Map extracted TTPs to MITRE ATT&CK technique IDs."""

    def __init__(self):
        self.techniques = {}
        self.tactics = {}
        self._loaded = False
        self._load_attack_data()

    def _load_attack_data(self):
        """Load ATT&CK data and verify integrity."""
        if not ATTACK_DATA_PATH.exists():
            logger.warning("MITRE ATT&CK data not found at %s", ATTACK_DATA_PATH)
            return

        try:
            with open(ATTACK_DATA_PATH) as f:
                data = json.load(f)

            # Load tactics
            for tactic in data.get('tactics', []):
                tid = tactic['id']
                self.tactics[tid] = tactic['name']
                self.tactics[tactic['name'].lower()] = tid

            # Load techniques (if populated)
            for technique in data.get('techniques', []):
                tid = technique.get('id', '')
                name = technique.get('name', '')
                if tid and name:
                    self.techniques[tid] = {
                        'name': name,
                        'tactic': technique.get('tactic', ''),
                        'description': technique.get('description', ''),
                    }
                    self.techniques[name.lower()] = tid

            self._loaded = True
            logger.debug("Loaded %d tactics, %d techniques",
                        len(data.get('tactics', [])),
                        len(data.get('techniques', [])))

        except Exception as e:
            logger.error("Failed to load ATT&CK data: %s", e)

    def map_technique(self, technique_name: str, category: str = "") -> Optional[str]:
        """Map a technique name or description to a MITRE ATT&CK ID.

        Returns ATT&CK ID (e.g., 'T1566') or None if no match.
        """
        if not self._loaded:
            return None

        # Direct ID match
        if technique_name.upper().startswith('T') and technique_name[1:].replace('.', '').isdigit():
            if technique_name.upper() in self.techniques:
                return technique_name.upper()

        # Name match
        lower_name = technique_name.lower().strip()
        if lower_name in self.techniques:
            result = self.techniques[lower_name]
            if isinstance(result, str):
                return result

        # Category to tactic mapping
        category_tactic_map = {
            'initial_access': 'TA0001',
            'execution': 'TA0002',
            'persistence': 'TA0003',
            'privilege_escalation': 'TA0004',
            'defense_evasion': 'TA0005',
            'credential_access': 'TA0006',
            'discovery': 'TA0007',
            'lateral_movement': 'TA0008',
            'collection': 'TA0009',
            'exfiltration': 'TA0010',
            'command_and_control': 'TA0011',
            'impact': 'TA0040',
            'resource_development': 'TA0042',
            'reconnaissance': 'TA0043',
            # Aliases
            'c2': 'TA0011',
            'exfil': 'TA0010',
            'privesc': 'TA0004',
            'recon': 'TA0043',
            'surveillance': 'TA0009',
        }

        return category_tactic_map.get(category.lower().strip())

    def get_tactic_name(self, tactic_id: str) -> Optional[str]:
        """Get tactic name from ID."""
        return self.tactics.get(tactic_id)

    def get_technique_info(self, technique_id: str) -> Optional[dict]:
        """Get technique details from ID."""
        return self.techniques.get(technique_id) if isinstance(
            self.techniques.get(technique_id), dict) else None

    def enrich_ttp(self, technique: str, category: str) -> dict:
        """Enrich a TTP with MITRE mapping.

        Returns dict with mitre_id and tactic if found.
        """
        result = {'mitre_id': None, 'tactic': None}

        mitre_id = self.map_technique(technique, category)
        if mitre_id:
            if mitre_id.startswith('TA'):
                result['tactic'] = mitre_id
                result['tactic_name'] = self.get_tactic_name(mitre_id)
            else:
                result['mitre_id'] = mitre_id
                info = self.get_technique_info(mitre_id)
                if info:
                    result['tactic'] = info.get('tactic')

        return result