"""
APK Deep Diff Analyzer - DEX File Analyzer
"""
import struct
import os
import zipfile
from typing import Dict, List, Any, Optional, Set
import re


class DexAnalyzer:
    """Analyze DEX (Dalvik Executable) files."""

    DEX_MAGIC = b'dex\n'

    def __init__(self, dex_path: str):
        self.dex_path = dex_path
        self.data = None
        self._load()

    def _load(self):
        """Load DEX file."""
        try:
            with open(self.dex_path, 'rb') as f:
                self.data = f.read()
        except Exception:
            self.data = b''

    def is_valid_dex(self) -> bool:
        """Check if file is a valid DEX."""
        return self.data.startswith(self.DEX_MAGIC) if self.data else False

    def get_header_info(self) -> Dict[str, Any]:
        """Parse DEX header information."""
        if not self.is_valid_dex():
            return {}

        try:
            # DEX header format (simplified)
            magic = self.data[:8]
            checksum = struct.unpack('<I', self.data[8:12])[0]
            signature = self.data[12:32].hex()
            file_size = struct.unpack('<I', self.data[32:36])[0]
            header_size = struct.unpack('<I', self.data[36:40])[0]
            endian_tag = struct.unpack('<I', self.data[40:44])[0]

            # String IDs
            string_ids_size = struct.unpack('<I', self.data[56:60])[0]
            string_ids_off = struct.unpack('<I', self.data[60:64])[0]

            # Type IDs
            type_ids_size = struct.unpack('<I', self.data[64:68])[0]
            type_ids_off = struct.unpack('<I', self.data[68:72])[0]

            # Proto IDs
            proto_ids_size = struct.unpack('<I', self.data[72:76])[0]
            proto_ids_off = struct.unpack('<I', self.data[76:80])[0]

            # Field IDs
            field_ids_size = struct.unpack('<I', self.data[80:84])[0]
            field_ids_off = struct.unpack('<I', self.data[84:88])[0]

            # Method IDs
            method_ids_size = struct.unpack('<I', self.data[88:92])[0]
            method_ids_off = struct.unpack('<I', self.data[92:96])[0]

            # Class defs
            class_defs_size = struct.unpack('<I', self.data[96:100])[0]
            class_defs_off = struct.unpack('<I', self.data[100:104])[0]

            return {
                "magic": magic.hex(),
                "checksum": checksum,
                "signature": signature,
                "file_size": file_size,
                "header_size": header_size,
                "endian_tag": endian_tag,
                "string_ids_count": string_ids_size,
                "type_ids_count": type_ids_size,
                "proto_ids_count": proto_ids_size,
                "field_ids_count": field_ids_size,
                "method_ids_count": method_ids_size,
                "class_defs_count": class_defs_size,
            }
        except Exception as e:
            return {"error": str(e)}

    def extract_strings(self) -> List[str]:
        """Extract strings from DEX file."""
        if not self.is_valid_dex():
            return []

        strings = []
        try:
            # Simple string extraction - look for readable ASCII/UTF-8 strings
            # This is a basic approach; full DEX parsing is complex
            pattern = rb'[\x20-\x7E]{4,}'
            matches = re.findall(pattern, self.data)
            for match in matches:
                try:
                    s = match.decode('utf-8', errors='ignore')
                    if len(s) >= 4:
                        strings.append(s)
                except:
                    pass

            # Also try to extract from string data section
            string_ids_size = struct.unpack('<I', self.data[56:60])[0]
            string_ids_off = struct.unpack('<I', self.data[60:64])[0]

            for i in range(min(string_ids_size, 10000)):  # Limit for performance
                try:
                    offset = struct.unpack('<I', self.data[string_ids_off + i*4:string_ids_off + i*4 + 4])[0]
                    # Read MUTF-8 string
                    str_len, bytes_read = self._read_uleb128(offset)
                    if str_len > 0 and str_len < 10000:
                        str_data = self.data[offset + bytes_read:offset + bytes_read + str_len]
                        s = str_data.decode('utf-8', errors='ignore')
                        if s and len(s) > 1:
                            strings.append(s)
                except:
                    pass
        except Exception:
            pass

        return sorted(set(strings))

    def _read_uleb128(self, offset: int) -> tuple:
        """Read unsigned LEB128 value."""
        result = 0
        shift = 0
        bytes_read = 0
        while True:
            byte = self.data[offset + bytes_read]
            result |= (byte & 0x7f) << shift
            bytes_read += 1
            if (byte & 0x80) == 0:
                break
            shift += 7
        return result, bytes_read

    def extract_urls(self) -> List[str]:
        """Extract URLs from DEX strings."""
        strings = self.extract_strings()
        urls = []
        url_pattern = re.compile(
            r'https?://[^\s<>"{}|\^`\[\]]+',
            re.IGNORECASE
        )
        for s in strings:
            matches = url_pattern.findall(s)
            urls.extend(matches)
        return sorted(set(urls))

    def extract_domains(self) -> List[str]:
        """Extract domain names from strings."""
        strings = self.extract_strings()
        domains = []
        domain_pattern = re.compile(
            r'(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}',
            re.IGNORECASE
        )
        for s in strings:
            matches = domain_pattern.findall(s)
            domains.extend(matches)
        return sorted(set(domains))

    def extract_ips(self) -> List[str]:
        """Extract IP addresses from strings."""
        strings = self.extract_strings()
        ips = []
        ip_pattern = re.compile(
            r'\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\b'
        )
        for s in strings:
            matches = ip_pattern.findall(s)
            ips.extend(matches)
        return sorted(set(ips))

    def extract_api_calls(self) -> List[str]:
        """Extract suspicious API references."""
        strings = self.extract_strings()
        suspicious = []
        suspicious_patterns = [
            r'Landroid/telephony/SmsManager;',
            r'Landroid/content/pm/PackageManager;',
            r'Landroid/app/admin/DevicePolicyManager;',
            r'Ljava/net/HttpURLConnection;',
            r'Landroid/location/LocationManager;',
            r'Landroid/hardware/Camera;',
            r'getDeviceId',
            r'sendTextMessage',
            r'exec\(',
            r'Runtime\.getRuntime',
            r'DexClassLoader',
            r'PathClassLoader',
            r'loadClass',
            r'invoke',
            r'getMethod',
            r'getDeclaredField',
        ]
        for s in strings:
            for pattern in suspicious_patterns:
                if re.search(pattern, s, re.IGNORECASE):
                    suspicious.append(s)
                    break
        return sorted(set(suspicious))

    def get_full_analysis(self) -> Dict[str, Any]:
        """Get complete DEX analysis."""
        header = self.get_header_info()
        return {
            "valid": self.is_valid_dex(),
            "header": header,
            "strings_count": len(self.extract_strings()),
            "urls": self.extract_urls(),
            "domains": self.extract_domains(),
            "ips": self.extract_ips(),
            "suspicious_apis": self.extract_api_calls(),
        }


def find_dex_files(extract_dir: str) -> List[str]:
    """Find all DEX files in extracted APK directory."""
    dex_files = []
    for root, dirs, files in os.walk(extract_dir):
        for f in files:
            if f.endswith('.dex'):
                dex_files.append(os.path.join(root, f))
    # Sort by name (classes.dex, classes2.dex, etc.)
    dex_files.sort(key=lambda x: (
        0 if os.path.basename(x) == 'classes.dex' else 1,
        os.path.basename(x)
    ))
    return dex_files


def analyze_all_dex(extract_dir: str) -> Dict[str, Any]:
    """Analyze all DEX files in directory."""
    dex_files = find_dex_files(extract_dir)
    results = {
        "dex_count": len(dex_files),
        "dex_files": {},
        "all_strings": set(),
        "all_urls": set(),
        "all_domains": set(),
        "all_ips": set(),
        "all_suspicious": set(),
    }

    for dex_path in dex_files:
        analyzer = DexAnalyzer(dex_path)
        filename = os.path.basename(dex_path)
        results["dex_files"][filename] = analyzer.get_full_analysis()

        if analyzer.is_valid_dex():
            results["all_strings"].update(analyzer.extract_strings())
            results["all_urls"].update(analyzer.extract_urls())
            results["all_domains"].update(analyzer.extract_domains())
            results["all_ips"].update(analyzer.extract_ips())
            results["all_suspicious"].update(analyzer.extract_api_calls())

    # Convert sets to sorted lists
    results["all_strings"] = sorted(results["all_strings"])
    results["all_urls"] = sorted(results["all_urls"])
    results["all_domains"] = sorted(results["all_domains"])
    results["all_ips"] = sorted(results["all_ips"])
    results["all_suspicious"] = sorted(results["all_suspicious"])

    return results