"""HW6 library rules; inspect student source without importing it.

The general shared-state checker is passed in by the embedded Rust wrapper so
its import resolution is shared without changing the policy of other homework.
This checks common compression APIs, not arbitrary compression algorithms.
"""

import ast
from pathlib import Path
import sys

COMPRESSION_APIS = {
    "zlib": {"compress", "compressobj", "decompress", "decompressobj"},
    "gzip": {"compress", "decompress", "open", "GzipFile"},
    "bz2": {"compress", "decompress", "open", "BZ2File", "BZ2Compressor", "BZ2Decompressor"},
    "lzma": {"compress", "decompress", "open", "LZMAFile", "LZMACompressor", "LZMADecompressor"},
}


def _check_file(tree, current, root, checker):
    violations = set()
    pending = checker._local_imports(tree, current, root)
    modules, functions = checker._dynamic_import_names(tree)
    strings = checker._static_string_values(tree)
    symbols = {}

    def report(node, message):
        violations.add(f"{current}:{node.lineno}:{node.col_offset + 1}: {message}")

    def check_module(name, node, base=root):
        parts = name.split(".")
        local = checker._candidate_module(base, parts)
        # Inspect package initializers as well as the explicitly imported file.
        for index in range(1, len(parts) + 1):
            candidate = checker._candidate_module(base, parts[:index])
            if candidate is not None:
                pending.add(candidate)
        namespace = base.joinpath(*parts).resolve()
        local_namespace = namespace.is_relative_to(root) and namespace.is_dir()
        if local is None and not local_namespace and parts[0] not in sys.stdlib_module_names and parts[0] != "anysystem":
            report(node, f"library '{name}' is not allowed: use AnySystem and the Python standard library")

    def dynamic_name(node):
        argument = node.args[0] if node.args else next(
            (item.value for item in node.keywords if item.arg == "name"), None
        )
        return checker._static_string(argument, strings) if argument is not None else None

    def symbol(node):
        if isinstance(node, ast.Name):
            return symbols.get(node.id)
        if isinstance(node, ast.Call) and checker._is_dynamic_import_call(node, modules, functions):
            return dynamic_name(node)
        resolved = checker._resolved_attribute(node)
        if resolved is not None:
            parent = symbol(resolved[0])
            if parent is not None:
                return f"{parent}.{resolved[1]}"
        return None

    def check_compression(name, node):
        parts = name.split(".")
        if len(parts) >= 2 and parts[1] in COMPRESSION_APIS.get(parts[0], set()):
            report(node, f"message compression is forbidden in HW6: '{parts[0]}.{parts[1]}'")

    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                check_module(alias.name, node)
                symbols[alias.asname or alias.name.split(".")[0]] = alias.name if alias.asname else alias.name.split(".")[0]
        elif isinstance(node, ast.ImportFrom):
            base = root
            if node.level:
                base = current.parent
                for _ in range(node.level - 1):
                    base = base.parent
                if not base.is_relative_to(root):
                    report(node, "local imports must stay inside the solution directory")
                    continue
            if node.module:
                check_module(node.module, node, base)
            for alias in node.names:
                name = f"{node.module}.{alias.name}" if node.module else alias.name
                if checker._candidate_module(base, name.split(".")) is not None:
                    check_module(name, node, base)
                if alias.name == "*" and node.module in COMPRESSION_APIS:
                    report(node, f"message compression is forbidden in HW6: star import from '{node.module}'")
                else:
                    check_compression(name, node)
                symbols[alias.asname or alias.name] = name
        elif isinstance(node, ast.Call) and checker._is_dynamic_import_call(node, modules, functions):
            name = dynamic_name(node)
            if name is None or name.startswith("."):
                report(node, "dynamic imports must use a statically known absolute module name")
            else:
                check_module(name, node)

    # Follow aliases, including compressor functions saved for later use.
    assignments = [parts for node in ast.walk(tree) if (parts := checker._assignment_parts(node)) is not None]
    for _ in range(len(assignments) + 1):
        changed = False
        for value, targets in assignments:
            name = symbol(value)
            if name is None:
                continue
            for target in targets:
                if isinstance(target, ast.Name) and target.id not in symbols:
                    symbols[target.id] = name
                    changed = True
        if not changed:
            break
    for node in ast.walk(tree):
        name = symbol(node)
        if name is not None:
            check_compression(name, node)
    return violations, pending


def validate_solution(path, checker):
    entrypoint = Path(path).resolve()
    root = entrypoint.parent
    pending = {entrypoint}
    visited = set()
    violations = set()
    while pending:
        current = pending.pop()
        if current in visited:
            continue
        visited.add(current)
        tree = ast.parse(current.read_text(encoding="utf-8-sig"), filename=str(current))
        errors, imports = _check_file(tree, current, root, checker)
        violations.update(errors)
        pending.update(imports - visited)
    return sorted(violations)