"""
Main pickle security scanner.

Orchestrates format detection, parsing, analysis, and reporting.
"""

from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional, Tuple, Union
import os

from pickle_scanner.core.opcodes import ParsedOpcode
from pickle_scanner.core.stack_machine import (
    StackMachine, GlobalRef, ReduceCall, BuildCall, StackItem, ItemType,
)
from pickle_scanner.formats.detector import (
    detect_format,
    detect_format_from_bytes,
    FileFormat,
    FormatInfo,
    is_safe_format,
)
from pickle_scanner.formats.pickle_parser import PickleParser, PickleStream
from pickle_scanner.formats.torch_parser import TorchParser, TorchArchive
from pickle_scanner.formats.joblib_parser import JoblibParser
from pickle_scanner.analysis.threat_intel import (
    ThreatDatabase,
    get_threat_database,
    is_dangerous_callable,
    RiskLevel as ThreatRiskLevel,
)
from pickle_scanner.analysis.obfuscation import (
    ObfuscationDetector,
    ObfuscationFinding,
    detect_obfuscation,
)
from pickle_scanner.analysis.ml_patterns import (
    MLPatternMatcher,
    detect_ml_framework,
    is_safe_ml_pattern,
    MLFramework,
)
from pickle_scanner.reporting.report import (
    ScanReport,
    Finding,
    ThreatScore,
    GlobalInfo,
    RiskLevel,
    Severity,
)


@dataclass
class ScanResult:
    """Result of a pickle scan."""
    report: ScanReport
    success: bool = True
    error: Optional[str] = None


@dataclass
class ScannerConfig:
    """Configuration for the scanner."""
    # Analysis options
    deep_analysis: bool = True      # Enable stack machine analysis
    detect_obfuscation: bool = True # Check for obfuscation techniques
    ml_pattern_matching: bool = True # Recognize safe ML patterns

    # Scoring adjustments
    ml_context_bonus: int = -15     # Score reduction for recognized ML patterns
    obfuscation_penalty_max: int = 30  # Maximum obfuscation penalty

    # Output options
    verbose: bool = False
    include_safe_patterns: bool = False  # Include safe pattern matches in report


class PickleScanner:
    """Main pickle security scanner."""

    def __init__(self, config: Optional[ScannerConfig] = None):
        self.config = config or ScannerConfig()
        self.threat_db = get_threat_database()
        self._last_error: Optional[str] = None

    def scan_file(self, path: Union[str, Path]) -> ScanResult:
        """Scan a file for pickle-based threats."""
        path = Path(path)

        if not path.exists():
            return ScanResult(
                report=ScanReport(file_path=str(path), format="unknown"),
                success=False,
                error=f"File not found: {path}"
            )

        try:
            file_size = path.stat().st_size
            with open(path, 'rb') as f:
                data = f.read()

            result = self.scan_bytes(data, str(path))
            result.report.file_size = file_size
            return result

        except Exception as e:
            return ScanResult(
                report=ScanReport(file_path=str(path), format="unknown"),
                success=False,
                error=f"Error reading file: {e}"
            )

    def scan_bytes(self, data: bytes, source: str = "<bytes>") -> ScanResult:
        """Scan bytes for pickle-based threats."""
        # Detect format, passing file extension for format-specific detection
        suffix = Path(source).suffix if source != "<bytes>" else ""
        format_info = detect_format_from_bytes(data, suffix)

        # Create report
        report = ScanReport(
            file_path=source,
            format=format_info.format.name.lower(),
            file_size=len(data),
        )

        # Safe formats don't need scanning
        if is_safe_format(format_info.format):
            report.threat_score = ThreatScore(
                base_score=0,
                confidence=format_info.confidence,
            )
            return ScanResult(report=report)

        # Route to appropriate parser
        try:
            if format_info.format == FileFormat.TORCHSERVE_MAR:
                self._scan_mar(data, report)
            elif format_info.format == FileFormat.NEMO:
                self._scan_nemo(data, report)
            elif format_info.format == FileFormat.JOBLIB:
                self._scan_joblib(data, report)
            elif format_info.format == FileFormat.KERAS:
                self._scan_keras(data, report)
            elif format_info.format == FileFormat.PYTORCH_ZIP:
                self._scan_pytorch(data, report, format_info)
            elif format_info.format == FileFormat.PICKLE:
                self._scan_raw_pickle(data, report)
            elif format_info.format == FileFormat.NUMPY_NPZ:
                self._scan_npz(data, report)
            else:
                # Try as raw pickle
                self._scan_raw_pickle(data, report)

        except Exception as e:
            report.errors.append(f"Analysis error: {e}")
            self._last_error = str(e)
            # Fail-closed: analysis errors may indicate EOP attack
            report.threat_score.base_score = max(
                report.threat_score.base_score, 50
            )
            report.add_finding(Finding(
                rule="eop_analysis_error",
                severity=Severity.MEDIUM,
                position=0,
                message=f"Analysis error may indicate EOP attack: {e}",
                evidence=str(e),
                remediation="Inspect file manually; analysis could not complete",
            ))

        return ScanResult(report=report)

    def _scan_pytorch(
        self,
        data: bytes,
        report: ScanReport,
        format_info: FormatInfo
    ) -> None:
        """Scan a PyTorch ZIP container."""
        parser = TorchParser(data=data)
        archive = parser.parse()

        report.format = f"pytorch_{archive.format}"
        report.pickle_count = len(archive.pickle_entries)
        report.errors.extend(archive.errors)

        # Fail-closed: errors with no pickle entries is suspicious (EOP)
        if archive.errors and not archive.pickle_entries:
            report.threat_score.base_score = max(
                report.threat_score.base_score, 75
            )
            report.add_finding(Finding(
                rule="eop_parse_failure",
                severity=Severity.HIGH,
                position=0,
                message="Archive parse errors with no extractable content — "
                        "possible EOP evasion attack",
                evidence="; ".join(archive.errors),
                remediation="File could not be fully parsed; treat as suspicious",
            ))

        # Analyze each pickle entry
        for entry in archive.pickle_entries:
            if entry.stream:
                self._analyze_pickle_stream(entry.stream, report)

        # Finalize scoring
        self._finalize_report(report)

    def _scan_raw_pickle(self, data: bytes, report: ScanReport) -> None:
        """Scan raw pickle data, handling multiple concatenated streams.

        PyTorch legacy format concatenates multiple pickle streams:
        stream 0: magic number, stream 1: version, stream 2: sys_info dict,
        stream 3: model state_dict, stream 4: storage keys.
        An attacker can inject a malicious stream anywhere in the sequence.
        """
        offset = 0
        stream_index = 0
        max_streams = 20  # Safety limit

        while offset < len(data) and stream_index < max_streams:
            remaining = data[offset:]
            if not remaining:
                break

            # Check if the remaining data starts with a valid pickle
            from pickle_scanner.formats.detector import PICKLE_PROTO_MAGIC
            is_pickle = False
            if remaining[0:1] == PICKLE_PROTO_MAGIC and len(remaining) > 1:
                if 2 <= remaining[1] <= 5:
                    is_pickle = True
            elif remaining[0:1] in [bytes([b]) for b in b'(cCdFgIlLNpPRSVX}])']:
                is_pickle = True

            # For the first stream, always attempt parsing even if the header
            # doesn't look like pickle (to record errors for fail-closed)
            if not is_pickle and stream_index > 0:
                break

            parser = PickleParser(remaining)
            stream = parser.parse()

            if stream_index == 0:
                report.protocol = stream.protocol
                report.errors.extend(stream.errors)
            # For subsequent streams, only propagate errors if the stream
            # had actual findings (avoid noise from tensor data)

            self._analyze_pickle_stream(stream, report)

            # Find where this stream ended (STOP opcode position + 1)
            stop_pos = None
            for op in stream.opcodes:
                if op.name == 'STOP':
                    stop_pos = op.position + 1  # STOP is 1 byte
                    break

            if stop_pos is None:
                # No STOP found — stream consumed all remaining data
                break

            offset += stop_pos
            stream_index += 1

        if stream_index > 1:
            report.pickle_count = stream_index

        self._finalize_report(report)

    def _scan_npz(self, data: bytes, report: ScanReport) -> None:
        """Scan a NumPy npz file (can contain pickles)."""
        import zipfile
        import io

        found_entries = False
        try:
            with zipfile.ZipFile(io.BytesIO(data)) as zf:
                for name in zf.namelist():
                    if name.endswith('.npy'):
                        found_entries = True
                        npy_data = zf.read(name)
                        # Check if npy uses pickle (object arrays)
                        if b'\x93NUMPY' in npy_data[:10]:
                            # NumPy format, check for object dtype
                            # Object arrays use pickle internally
                            if b"'O'" in npy_data[:200] or b"object" in npy_data[:200]:
                                report.errors.append(f"Object array in {name} may contain pickle")
        except Exception as e:
            report.errors.append(f"NPZ parse error: {e}")

        # Fail-closed: parse errors with no entries is suspicious
        if report.errors and not found_entries:
            report.threat_score.base_score = max(
                report.threat_score.base_score, 50
            )
            report.add_finding(Finding(
                rule="eop_npz_parse_failure",
                severity=Severity.MEDIUM,
                position=0,
                message="NPZ parse errors with no extractable content — "
                        "possible EOP evasion attack",
                evidence="; ".join(report.errors),
            ))

        self._finalize_report(report)

    def _scan_joblib(self, data: bytes, report: ScanReport) -> None:
        """Scan a Joblib file (raw or compressed pickle)."""
        parser = JoblibParser(data)
        result = parser.parse()

        report.format = f"joblib_{result.compression}"
        report.errors.extend(result.errors)

        if result.pickle_stream:
            report.protocol = result.pickle_stream.protocol
            report.errors.extend(result.pickle_stream.errors)
            self._analyze_pickle_stream(result.pickle_stream, report)

        # Fail-closed: errors with no pickle stream is suspicious
        if result.errors and not result.pickle_stream:
            report.threat_score.base_score = max(
                report.threat_score.base_score, 50
            )
            report.add_finding(Finding(
                rule="eop_joblib_parse_failure",
                severity=Severity.MEDIUM,
                position=0,
                message="Joblib parse error with no extractable pickle — "
                        "possible EOP evasion",
                evidence="; ".join(result.errors),
            ))

        self._finalize_report(report)

    def _scan_keras(self, data: bytes, report: ScanReport,
                    _depth: int = 0) -> None:
        """Scan a Keras .keras file (ZIP with npz weights)."""
        import zipfile
        import io

        MAX_DEPTH = 3
        if _depth >= MAX_DEPTH:
            report.errors.append("Max recursion depth reached scanning Keras file")
            self._finalize_report(report)
            return

        report.format = "keras"

        try:
            with zipfile.ZipFile(io.BytesIO(data)) as zf:
                for name in zf.namelist():
                    if name.endswith('.npz'):
                        # Nested NPZ — extract and scan
                        npz_data = zf.read(name)
                        try:
                            with zipfile.ZipFile(io.BytesIO(npz_data)) as npz_zf:
                                for npy_name in npz_zf.namelist():
                                    if npy_name.endswith('.npy'):
                                        npy_data = npz_zf.read(npy_name)
                                        if b'\x93NUMPY' in npy_data[:10]:
                                            if (b"'O'" in npy_data[:200] or
                                                    b"object" in npy_data[:200]):
                                                report.errors.append(
                                                    f"Object array in {name}/{npy_name} "
                                                    "may contain pickle"
                                                )
                                        # Also check for raw pickle inside npy
                                        pickle_info = self._try_pickle_scan(npy_data)
                                        if pickle_info:
                                            self._analyze_pickle_stream(
                                                pickle_info, report
                                            )
                        except zipfile.BadZipFile:
                            report.errors.append(f"Invalid NPZ inside keras: {name}")

                    elif name.endswith('.pkl') or name.endswith('.pickle'):
                        # Direct pickle files inside .keras
                        pkl_data = zf.read(name)
                        parser = PickleParser(pkl_data)
                        stream = parser.parse()
                        self._analyze_pickle_stream(stream, report)

        except zipfile.BadZipFile as e:
            report.errors.append(f"Invalid Keras ZIP: {e}")
            report.threat_score.base_score = max(
                report.threat_score.base_score, 50
            )
            report.add_finding(Finding(
                rule="eop_keras_parse_failure",
                severity=Severity.MEDIUM,
                position=0,
                message="Keras file parse error — possible EOP evasion",
                evidence=str(e),
            ))

        self._finalize_report(report)

    def _scan_mar(self, data: bytes, report: ScanReport) -> None:
        """Scan a TorchServe .mar file (ZIP with nested model archives).

        MAR files are ZIP archives containing:
        - MAR-INF/MANIFEST.json (metadata)
        - A serialized model file (.pt, .pth, or another archive)
        - Handler code (Python files — not scanned here, but inner models are)
        """
        import zipfile
        import io

        report.format = "torchserve_mar"

        try:
            with zipfile.ZipFile(io.BytesIO(data)) as zf:
                for name in zf.namelist():
                    member_data = zf.read(name)

                    # Scan any pickle files directly
                    if name.endswith('.pkl') or name.endswith('.pickle'):
                        parser = PickleParser(member_data)
                        stream = parser.parse()
                        self._analyze_pickle_stream(stream, report)

                    # Scan nested ZIP archives (e.g., .pt/.pth inside .mar)
                    elif member_data[:4] in (b'PK\x03\x04', b'PK\x05\x06'):
                        inner_parser = TorchParser(data=member_data)
                        inner_archive = inner_parser.parse()
                        report.errors.extend(inner_archive.errors)
                        for entry in inner_archive.pickle_entries:
                            if entry.stream:
                                self._analyze_pickle_stream(
                                    entry.stream, report
                                )

                    # Scan nested tar archives
                    elif (len(member_data) >= 512 and
                          member_data[257:262] in (b'ustar', b'ustar ')) or \
                         member_data[:2] == b'\x1f\x8b':
                        inner_parser = TorchParser(data=member_data)
                        inner_archive = inner_parser.parse()
                        report.errors.extend(inner_archive.errors)
                        for entry in inner_archive.pickle_entries:
                            if entry.stream:
                                self._analyze_pickle_stream(
                                    entry.stream, report
                                )

        except zipfile.BadZipFile as e:
            report.errors.append(f"Invalid MAR ZIP: {e}")
            report.threat_score.base_score = max(
                report.threat_score.base_score, 50
            )
            report.add_finding(Finding(
                rule="eop_mar_parse_failure",
                severity=Severity.MEDIUM,
                position=0,
                message="MAR file parse error — possible EOP evasion",
                evidence=str(e),
            ))

        self._finalize_report(report)

    def _scan_nemo(self, data: bytes, report: ScanReport) -> None:
        """Scan a NeMo .nemo file (tar archive with torch/joblib models).

        NeMo files are tar archives (often gzipped) containing:
        - .pt/.pth files (PyTorch models via torch.save)
        - .pkl/.pickle files (raw pickle)
        - Joblib-compressed pickle files
        - model_config.yaml (metadata)
        """
        import tarfile
        import io

        report.format = "nemo"

        try:
            with tarfile.open(fileobj=io.BytesIO(data), mode='r:*') as tf:
                for member in tf.getmembers():
                    if not member.isfile():
                        continue

                    try:
                        f = tf.extractfile(member)
                        if not f:
                            continue
                        member_data = f.read()
                    except Exception as e:
                        report.errors.append(
                            f"Error reading {member.name}: {e}"
                        )
                        continue

                    name = member.name

                    # Direct pickle files
                    if name.endswith('.pkl') or name.endswith('.pickle'):
                        parser = PickleParser(member_data)
                        stream = parser.parse()
                        self._analyze_pickle_stream(stream, report)

                    # PyTorch model files (ZIP-based)
                    elif (name.endswith('.pt') or name.endswith('.pth')) and \
                         member_data[:4] in (b'PK\x03\x04', b'PK\x05\x06'):
                        inner_parser = TorchParser(data=member_data)
                        inner_archive = inner_parser.parse()
                        report.errors.extend(inner_archive.errors)
                        for entry in inner_archive.pickle_entries:
                            if entry.stream:
                                self._analyze_pickle_stream(
                                    entry.stream, report
                                )

                    # Joblib-compressed files
                    elif name.endswith('.joblib'):
                        jparser = JoblibParser(member_data)
                        jresult = jparser.parse()
                        report.errors.extend(jresult.errors)
                        if jresult.pickle_stream:
                            self._analyze_pickle_stream(
                                jresult.pickle_stream, report
                            )

                    # Check for compressed pickle (joblib without extension)
                    else:
                        from pickle_scanner.formats.joblib_parser import (
                            detect_joblib_compression,
                        )
                        comp = detect_joblib_compression(member_data)
                        if comp != 'none':
                            jparser = JoblibParser(member_data)
                            jresult = jparser.parse()
                            report.errors.extend(jresult.errors)
                            if jresult.pickle_stream:
                                self._analyze_pickle_stream(
                                    jresult.pickle_stream, report
                                )

        except tarfile.TarError as e:
            report.errors.append(f"Invalid NeMo tar: {e}")
            report.threat_score.base_score = max(
                report.threat_score.base_score, 50
            )
            report.add_finding(Finding(
                rule="eop_nemo_parse_failure",
                severity=Severity.MEDIUM,
                position=0,
                message="NeMo file parse error — possible EOP evasion",
                evidence=str(e),
            ))

        self._finalize_report(report)

    def _try_pickle_scan(self, data: bytes) -> Optional[PickleStream]:
        """Try to parse data as pickle, return stream or None."""
        from pickle_scanner.formats.detector import _detect_pickle
        if _detect_pickle(data):
            try:
                parser = PickleParser(data)
                return parser.parse()
            except Exception:
                pass
        return None

    def _analyze_pickle_stream(self, stream: PickleStream, report: ScanReport) -> None:
        """Analyze a parsed pickle stream."""
        # Run stack machine simulation if enabled
        machine = None
        if self.config.deep_analysis:
            machine = StackMachine()
            machine.protocol = stream.protocol
            for opcode in stream.opcodes:
                machine.process_opcode(opcode)
            report.errors.extend(machine.errors)

        # Analyze globals
        if machine:
            self._analyze_globals(machine.globals_found, report)

            # Analyze reduce calls for argument extraction
            for call in machine.reduce_calls:
                if call.callable_ref:
                    self._check_reduce_call(call, report)

            # Dataflow analysis: detect chained calls and suspicious arguments
            self._analyze_dataflow(machine, report)

        else:
            # Fallback: extract globals directly from opcodes
            self._analyze_globals_from_opcodes(stream.opcodes, report)

        # Detect obfuscation
        if self.config.detect_obfuscation:
            obfuscation_findings = detect_obfuscation(stream.opcodes, machine)
            for finding in obfuscation_findings:
                report.obfuscation_techniques.append(finding.technique.name)
                report.threat_score.obfuscation_penalty += finding.severity

                # Add finding
                report.add_finding(Finding(
                    rule=f"obfuscation_{finding.technique.name.lower()}",
                    severity=Severity.HIGH if finding.severity >= 20 else Severity.MEDIUM,
                    position=finding.position,
                    message=finding.description,
                    evidence=finding.evidence,
                ))

        # ML framework detection
        if self.config.ml_pattern_matching and machine:
            framework = detect_ml_framework(machine.globals_found)
            if framework != MLFramework.UNKNOWN:
                report.ml_framework_detected = framework.name

    def _analyze_globals(
        self,
        globals_found: List[Tuple[GlobalRef, int]],
        report: ScanReport
    ) -> None:
        """Analyze global references for threats."""
        ml_matcher = MLPatternMatcher() if self.config.ml_pattern_matching else None
        dangerous_count = 0
        safe_ml_count = 0

        for ref, position in globals_found:
            # Check if dangerous
            is_dangerous, risk_info = is_dangerous_callable(ref.module, ref.name)

            # Check if safe ML pattern
            is_safe_ml = False
            ml_desc = None
            if ml_matcher:
                match = ml_matcher.match_global(ref)
                if match:
                    is_safe_ml = True
                    ml_desc = match.description
                    safe_ml_count += 1

            # Create global info
            global_info = GlobalInfo(
                module=ref.module,
                name=ref.name,
                position=position,
                opcode=ref.opcode,
                is_dangerous=is_dangerous,
                is_safe_ml=is_safe_ml,
                risk_description=risk_info.description if risk_info else ml_desc,
            )
            report.add_global(global_info)

            # Generate findings for dangerous callables
            if is_dangerous and risk_info:
                dangerous_count += 1

                # Determine severity
                if risk_info.risk_level >= ThreatRiskLevel.CRITICAL:
                    severity = Severity.CRITICAL
                elif risk_info.risk_level >= ThreatRiskLevel.HIGH:
                    severity = Severity.HIGH
                elif risk_info.risk_level >= ThreatRiskLevel.MEDIUM:
                    severity = Severity.MEDIUM
                else:
                    severity = Severity.LOW

                report.add_finding(Finding(
                    rule=f"dangerous_callable_{risk_info.category}",
                    severity=severity,
                    position=position,
                    message=f"Dangerous callable: {ref.full_path} - {risk_info.description}",
                    callable=ref.full_path,
                    opcode=ref.opcode,
                    evidence=f"{ref.opcode} {ref.module} {ref.name}",
                ))

                # Update base score
                report.threat_score.base_score = max(
                    report.threat_score.base_score,
                    risk_info.risk_level
                )

        # Apply ML context adjustment
        if safe_ml_count > 0 and dangerous_count == 0:
            report.threat_score.context_adjustment = self.config.ml_context_bonus

    def _analyze_globals_from_opcodes(
        self,
        opcodes: List[ParsedOpcode],
        report: ScanReport
    ) -> None:
        """Fallback: analyze globals directly from opcodes."""
        for op in opcodes:
            if op.name in ('GLOBAL', 'INST'):
                if isinstance(op.arg, tuple) and len(op.arg) == 2:
                    module, name = op.arg
                elif isinstance(op.arg, str) and '\n' in op.arg:
                    parts = op.arg.split('\n')
                    module, name = parts[0], parts[1] if len(parts) > 1 else ""
                else:
                    continue

                ref = GlobalRef(module, name, op.name)
                is_dangerous, risk_info = is_dangerous_callable(module, name)

                global_info = GlobalInfo(
                    module=module,
                    name=name,
                    position=op.position,
                    opcode=op.name,
                    is_dangerous=is_dangerous,
                    risk_description=risk_info.description if risk_info else None,
                )
                report.add_global(global_info)

                if is_dangerous and risk_info:
                    severity = Severity.CRITICAL if risk_info.risk_level >= ThreatRiskLevel.CRITICAL else Severity.HIGH

                    report.add_finding(Finding(
                        rule=f"dangerous_callable_{risk_info.category}",
                        severity=severity,
                        position=op.position,
                        message=f"Dangerous callable: {module}.{name}",
                        callable=f"{module}.{name}",
                        opcode=op.name,
                    ))

                    report.threat_score.base_score = max(
                        report.threat_score.base_score,
                        risk_info.risk_level
                    )

    def _check_reduce_call(self, call, report: ScanReport) -> None:
        """Check a REDUCE call for additional context."""
        ref = call.callable_ref
        if not ref:
            return

        # Try to extract argument information
        args_str = None
        if call.args is not None:
            try:
                args_str = repr(call.args)[:200]  # Limit length
            except:
                pass

        # Find existing finding and update with args
        for finding in report.findings:
            if finding.callable == ref.full_path and finding.position == call.position:
                finding.arguments = args_str
                break

    # Patterns that indicate code/shell content in string arguments
    _CODE_PATTERNS = [
        'import ', '__import__', 'os.system', 'os.popen', 'subprocess',
        'eval(', 'exec(', 'compile(', '__builtins__', '__class__',
        '__subclasses__', '__globals__', '__reduce__',
        '/bin/sh', '/bin/bash', 'cmd.exe', 'powershell',
        'echo ', 'curl ', 'wget ', 'chmod ', 'rm -',
    ]

    def _analyze_dataflow(self, machine: StackMachine,
                          report: ScanReport) -> None:
        """Analyze dataflow across reduce calls for chained attacks.

        Detects:
        1. Chained reduce calls: output of one REDUCE feeds another
        2. Unknown callables receiving suspicious string arguments
        3. Helper gadgets (getattr equivalents) used to resolve callables
        """
        ml_matcher = MLPatternMatcher() if self.config.ml_pattern_matching else None

        # Build a set of reduce call positions for chain detection
        reduce_positions = {call.position for call in machine.reduce_calls}

        # Build a set of known-dangerous global paths for quick lookup
        dangerous_globals = set()
        for ref, _ in machine.globals_found:
            is_dangerous, _ = is_dangerous_callable(ref.module, ref.name)
            if is_dangerous:
                dangerous_globals.add(ref.full_path)

        for call in machine.reduce_calls:
            ref = call.callable_ref
            if not ref:
                continue

            # --- Check 1: Chained reduce calls ---
            # If a reduce call's arguments contain another ReduceCall result,
            # this is a chained invocation pattern (common in gadget attacks)
            if self._args_contain_reduce(call.args):
                # Only flag if the callable itself is suspicious
                is_safe_ml = False
                if ml_matcher:
                    is_safe_ml = ml_matcher.match_global(ref) is not None

                if not is_safe_ml:
                    is_dangerous, risk_info = is_dangerous_callable(
                        ref.module, ref.name
                    )
                    # Flag if the callable is dangerous or if it chains into
                    # something dangerous
                    if is_dangerous or self._chain_reaches_dangerous(
                        call, machine.reduce_calls, dangerous_globals
                    ):
                        if not any(
                            f.callable == ref.full_path and
                            f.rule == "dataflow_chained_reduce"
                            for f in report.findings
                        ):
                            report.add_finding(Finding(
                                rule="dataflow_chained_reduce",
                                severity=Severity.HIGH,
                                position=call.position,
                                message=f"Chained reduce call: {ref.full_path} "
                                        "receives output of another call — "
                                        "possible gadget chain attack",
                                callable=ref.full_path,
                                evidence=self._summarize_args(call.args),
                            ))
                            report.threat_score.base_score = max(
                                report.threat_score.base_score, 75
                            )

            # --- Check 2: Unknown callable with suspicious string args ---
            is_known_dangerous, _ = is_dangerous_callable(ref.module, ref.name)
            is_safe_ml = False
            if ml_matcher:
                is_safe_ml = ml_matcher.match_global(ref) is not None

            if not is_known_dangerous and not is_safe_ml:
                # Check if string arguments look like code or commands
                suspicious_args = self._find_suspicious_args(call.args)
                if suspicious_args:
                    report.add_finding(Finding(
                        rule="dataflow_suspicious_args",
                        severity=Severity.HIGH,
                        position=call.position,
                        message=f"Unknown callable {ref.full_path} receives "
                                f"suspicious argument: {suspicious_args[:100]}",
                        callable=ref.full_path,
                        evidence=suspicious_args[:200],
                    ))
                    report.threat_score.base_score = max(
                        report.threat_score.base_score, 75
                    )

    def _args_contain_reduce(self, args: any) -> bool:
        """Check if arguments contain a ReduceCall result."""
        if isinstance(args, ReduceCall):
            return True
        if isinstance(args, (tuple, list)):
            return any(self._args_contain_reduce(a) for a in args)
        return False

    def _chain_reaches_dangerous(
        self,
        call: ReduceCall,
        all_calls: list,
        dangerous_globals: set,
    ) -> bool:
        """Check if a reduce call chain eventually reaches a dangerous callable."""
        if call.callable_ref and call.callable_ref.full_path in dangerous_globals:
            return True
        # Check if any downstream call in the chain is dangerous
        # (calls that come after this one in position order)
        for other in all_calls:
            if other.position > call.position and other.callable_ref:
                if other.callable_ref.full_path in dangerous_globals:
                    return True
        return False

    def _find_suspicious_args(self, args: any) -> Optional[str]:
        """Check if arguments contain strings that look like code/commands."""
        strings = self._extract_strings(args)
        for s in strings:
            s_lower = s.lower()
            for pattern in self._CODE_PATTERNS:
                if pattern.lower() in s_lower:
                    return s
        return None

    @staticmethod
    def _extract_strings(obj: any) -> list:
        """Recursively extract string values from arguments."""
        strings = []
        if isinstance(obj, str):
            strings.append(obj)
        elif isinstance(obj, (tuple, list)):
            for item in obj:
                strings.extend(PickleScanner._extract_strings(item))
        elif isinstance(obj, dict):
            for v in obj.values():
                strings.extend(PickleScanner._extract_strings(v))
        return strings

    @staticmethod
    def _summarize_args(args: any) -> str:
        """Create a short summary of call arguments for evidence."""
        try:
            return repr(args)[:200]
        except Exception:
            return "<unprintable>"

    def _finalize_report(self, report: ScanReport) -> None:
        """Finalize the report with scoring adjustments."""
        # Escalate for heavy obfuscation: many STACK_GLOBAL dynamic calls
        # combined with encoded payloads strongly indicate intentional evasion
        stack_global_count = sum(
            1 for f in report.findings
            if f.rule == "obfuscation_stack_global_dynamic"
        )
        has_encoded_payload = any(
            f.rule == "obfuscation_encoded_payload"
            for f in report.findings
        )
        # Only escalate if the file doesn't look like a legitimate ML model.
        # Legitimate frameworks (flair, PyTorch) use STACK_GLOBAL in protocol 4+
        # and may contain high-entropy strings (alphabets, embeddings) that
        # trigger encoded_payload detection. When an ML framework is detected
        # and no globals are dangerous, the obfuscation is benign.
        ml_framework_detected = report.ml_framework_detected is not None
        has_dangerous_globals = any(
            g.is_dangerous for g in report.globals_observed
        )
        benign_ml_context = ml_framework_detected and not has_dangerous_globals
        if stack_global_count >= 4 and has_encoded_payload and not benign_ml_context:
            report.threat_score.base_score = max(
                report.threat_score.base_score, 60
            )
            if not any(f.rule == "heavy_obfuscation" for f in report.findings):
                report.add_finding(Finding(
                    rule="heavy_obfuscation",
                    severity=Severity.HIGH,
                    position=0,
                    message=f"Heavy obfuscation detected: {stack_global_count} "
                            "dynamic STACK_GLOBAL calls with encoded payload "
                            "— likely evasion attempt",
                    evidence=f"{stack_global_count} STACK_GLOBAL + encoded payload",
                    remediation="File uses extensive obfuscation to hide callables; "
                                "treat as suspicious",
                ))
        # Cap obfuscation penalty
        report.threat_score.obfuscation_penalty = min(
            report.threat_score.obfuscation_penalty,
            self.config.obfuscation_penalty_max
        )

        # Calculate confidence based on analysis completeness
        confidence = 1.0
        if report.errors:
            confidence -= 0.1 * min(len(report.errors), 3)
        if report.format == "unknown":
            confidence -= 0.2

        report.threat_score.confidence = max(0.5, confidence)

        # Deduplicate obfuscation techniques
        report.obfuscation_techniques = list(set(report.obfuscation_techniques))


def scan_file(path: Union[str, Path], config: Optional[ScannerConfig] = None) -> ScanResult:
    """Convenience function to scan a file."""
    scanner = PickleScanner(config)
    return scanner.scan_file(path)


def scan_bytes(data: bytes, source: str = "<bytes>", config: Optional[ScannerConfig] = None) -> ScanResult:
    """Convenience function to scan bytes."""
    scanner = PickleScanner(config)
    return scanner.scan_bytes(data, source)
