import os
from typing import ClassVar

from .base import BaseScanner, IssueSeverity, ScanResult

try:
    import tflite

    HAS_TFLITE = True
except Exception:  # pragma: no cover - optional dependency
    HAS_TFLITE = False


class TFLiteScanner(BaseScanner):
    """Scanner for TensorFlow Lite model files."""

    name = "tflite"
    description = "Scans TensorFlow Lite models for integrity and safety issues"
    supported_extensions: ClassVar[list[str]] = [".tflite"]

    @classmethod
    def can_handle(cls, path: str) -> bool:
        if not HAS_TFLITE:
            return False
        return os.path.isfile(path) and os.path.splitext(path)[1].lower() in cls.supported_extensions

    def scan(self, path: str) -> ScanResult:
        path_check_result = self._check_path(path)
        if path_check_result:
            return path_check_result

        result = self._create_result()
        result.metadata["file_size"] = self.get_file_size(path)

        if not HAS_TFLITE:
            result.add_check(
                name="TFLite Library Check",
                passed=False,
                message="tflite package not installed. Install with 'pip install modelaudit[tflite]'",
                severity=IssueSeverity.WARNING,
                location=path,
                details={"required_package": "tflite"},
            )
            result.finish(success=False)
            return result

        try:
            with open(path, "rb") as f:
                data = f.read()
                result.bytes_scanned = len(data)

            # Check for TFLite magic bytes "TFL3" at offset 4
            # TFLite uses FlatBuffer format: bytes 0-3 are root table offset, bytes 4-7 are file identifier
            if len(data) < 8 or data[4:8] != b"TFL3":
                result.add_check(
                    name="TFLite Magic Bytes Check",
                    passed=False,
                    message="File does not have valid TFLite magic bytes (expected 'TFL3' at offset 4)",
                    severity=IssueSeverity.INFO,
                    location=path,
                    details={"magic_bytes_at_offset_4": data[4:8].hex() if len(data) >= 8 else "file_too_short"},
                    why="Valid TFLite files use FlatBuffer format with 'TFL3' identifier at bytes 4-7. "
                    "Missing or incorrect identifier may indicate file corruption or spoofing.",
                )
                result.finish(success=False)
                return result

            model = tflite.Model.GetRootAsModel(data, 0)
        except Exception as e:  # pragma: no cover - parse errors
            result.add_check(
                name="TFLite File Parse",
                passed=False,
                message=f"Invalid TFLite file or parse error: {e}",
                severity=IssueSeverity.INFO,
                location=path,
                details={"exception": str(e), "exception_type": type(e).__name__},
            )
            result.finish(success=False)
            return result

        subgraph_count = model.SubgraphsLength()
        result.metadata["subgraph_count"] = subgraph_count

        for sg_index in range(subgraph_count):
            subgraph = model.Subgraphs(sg_index)
            tensors_len = subgraph.TensorsLength()
            operators_len = subgraph.OperatorsLength()
            result.metadata.setdefault("tensor_counts", []).append(tensors_len)
            result.metadata.setdefault("operator_counts", []).append(operators_len)

            for o_index in range(operators_len):
                op = subgraph.Operators(o_index)
                opcode = model.OperatorCodes(op.OpcodeIndex())
                builtin = opcode.BuiltinCode()
                if builtin == tflite.BuiltinOperator.CUSTOM:
                    custom = opcode.CustomCode()
                    name = custom.decode("utf-8", "ignore") if custom else "unknown"
                    result.add_check(
                        name="Custom Operator Detection",
                        passed=False,
                        message=f"Model uses custom operator '{name}'",
                        severity=IssueSeverity.CRITICAL,
                        location=f"{path} (operator {o_index})",
                        details={"operator_name": name, "operator_index": o_index},
                    )

        result.finish(success=not result.has_errors)
        return result
