import traceback

import structlog
from temporalio.worker.workflow_sandbox import RestrictedWorkflowAccessError, SandboxRestrictions

logger = structlog.get_logger(__name__)

SANDBOX_PASSTHROUGH_MODULES = (
    "mistralai.workflows",
    "mistralai.client.models",
)

SANDBOX_RESTRICTIONS = SandboxRestrictions.default.with_passthrough_modules(*SANDBOX_PASSTHROUGH_MODULES)

_SANDBOX_INTERNAL_MARKERS = (
    "temporalio/",
    "<frozen importlib",
    "<string>",
)

_THIRD_PARTY_MARKERS = (
    "site-packages/",
    ".venv/",
)


def _find_sandbox_restriction_error(exc: BaseException) -> RestrictedWorkflowAccessError | None:
    """Walk the exception chain to find a RestrictedWorkflowAccessError."""
    current: BaseException | None = exc
    while current is not None:
        if isinstance(current, RestrictedWorkflowAccessError):
            return current
        current = current.__cause__
    return None


def _format_sandbox_import_chain(exc: RestrictedWorkflowAccessError) -> str:
    """Build a human-readable import chain from the sandbox error traceback.

    Categorizes frames into "your code" vs "third-party library" and filters
    out temporal/sandbox internals, so the user can see exactly which import
    in their code triggers the violation and which library is responsible.
    """
    tb = traceback.extract_tb(exc.__traceback__)

    # Filter out sandbox/temporal internals
    relevant_frames = [frame for frame in tb if not any(m in frame.filename for m in _SANDBOX_INTERNAL_MARKERS)]
    if not relevant_frames:
        return ""

    lines: list[str] = ["  Import chain:"]
    for frame in relevant_frames:
        is_third_party = any(m in frame.filename for m in _THIRD_PARTY_MARKERS)
        label = "[third-party]" if is_third_party else "[your code]  "
        lines.append(f"    {label} {frame.filename}:{frame.lineno}: {frame.line}")
    return "\n".join(lines)


_SANDBOX_ERROR_GUIDANCE = (
    "This usually means your workflow code imports or uses a module "
    "that is not allowed inside the sandbox "
    "(e.g. file I/O, network calls, threading).\n"
    "To fix this, either:\n"
    "  - Move the import inside an activity instead of the workflow\n"
    "  - Use `with workflow.unsafe.imports_passed_through():` "
    "to bypass the sandbox for safe imports"
)


def log_if_sandbox_restriction_error(exc: BaseException, context: str) -> None:
    """If exc (or its cause chain) contains a RestrictedWorkflowAccessError,
    log a structured error with the import chain and actionable guidance."""
    sandbox_error = _find_sandbox_restriction_error(exc)
    if sandbox_error is None:
        return
    import_chain = _format_sandbox_import_chain(sandbox_error)
    parts = [
        f"Workflow sandbox {context} failed: {sandbox_error}",
        import_chain,
        _SANDBOX_ERROR_GUIDANCE,
    ]
    message = "\n".join(part for part in parts if part)
    logger.error(message)
