"""Safe and deterministic firmware extraction helpers."""

from __future__ import annotations

from collections import deque
from dataclasses import dataclass
import gzip
import hashlib
import io
import shutil
import tarfile
from pathlib import Path, PurePosixPath
from typing import Protocol
import zipfile

from sentra.models import ExtractedArtifact

MAX_EXTRACTED_FILES = 128
MAX_TOTAL_BYTES = 8 * 1024 * 1024
MAX_EXTRACTION_DEPTH = 2


@dataclass(frozen=True)
class ExtractionRequest:
    source_path: Path
    destination_root: Path
    max_files: int = MAX_EXTRACTED_FILES
    max_total_bytes: int = MAX_TOTAL_BYTES
    max_depth: int = MAX_EXTRACTION_DEPTH


@dataclass(frozen=True)
class DetectedFileType:
    name: str
    media_type: str
    extractable: bool = False


@dataclass
class ExtractionState:
    artifacts: list[ExtractedArtifact]
    total_bytes: int = 0
    truncated: bool = False
    depth_limit_hit: bool = False
    nested_archives_seen: int = 0
    extraction_notes: list[str] | None = None

    def add_note(self, note: str) -> None:
        if self.extraction_notes is None:
            self.extraction_notes = []
        if note not in self.extraction_notes:
            self.extraction_notes.append(note)


@dataclass(frozen=True)
class ExtractionTask:
    display_name: str
    payload: bytes
    relative_root: Path
    depth: int
    extracted_from: str | None = None


@dataclass(frozen=True)
class ExtractionResult:
    provider: str
    extracted_root: Path | None
    artifacts: list[ExtractedArtifact]
    metadata: dict[str, object]


class ExtractionProvider(Protocol):
    name: str

    def supports(self, path: Path, payload: bytes) -> bool:
        """Return True when this provider can safely unpack the file."""

    def detected_type(self, path: Path, payload: bytes) -> DetectedFileType:
        """Describe the detected file type for this provider."""

    def extract_members(self, task: ExtractionTask) -> list[tuple[Path, bytes]]:
        """Return safe archive members as relative paths plus payload bytes."""


class ZipExtractionProvider:
    name = "zip"

    def supports(self, path: Path, payload: bytes) -> bool:
        return zipfile.is_zipfile(io.BytesIO(payload)) or path.suffix.lower() == ".zip"

    def detected_type(self, path: Path, payload: bytes) -> DetectedFileType:
        return DetectedFileType(name="zip", media_type="application/zip", extractable=True)

    def extract_members(self, task: ExtractionTask) -> list[tuple[Path, bytes]]:
        members: list[tuple[Path, bytes]] = []
        with zipfile.ZipFile(io.BytesIO(task.payload)) as archive:
            for member in sorted(
                (member for member in archive.infolist() if not member.is_dir()),
                key=lambda item: item.filename,
            ):
                relative = _safe_member_path(member.filename)
                with archive.open(member) as source:
                    members.append((relative, source.read()))
        return members


class TarExtractionProvider:
    name = "tar"

    def supports(self, path: Path, payload: bytes) -> bool:
        try:
            with tarfile.open(fileobj=io.BytesIO(payload), mode="r:*"):
                return True
        except tarfile.TarError:
            if path.exists():
                try:
                    return tarfile.is_tarfile(path)
                except OSError:
                    return False
            return False

    def detected_type(self, path: Path, payload: bytes) -> DetectedFileType:
        return DetectedFileType(name="tar", media_type="application/x-tar", extractable=True)

    def extract_members(self, task: ExtractionTask) -> list[tuple[Path, bytes]]:
        members: list[tuple[Path, bytes]] = []
        with tarfile.open(fileobj=io.BytesIO(task.payload), mode="r:*") as archive:
            for member in sorted((member for member in archive.getmembers() if member.isfile()), key=lambda item: item.name):
                relative = _safe_member_path(member.name)
                extracted = archive.extractfile(member)
                if extracted is None:
                    continue
                with extracted:
                    members.append((relative, extracted.read()))
        return members


class GzipExtractionProvider:
    name = "gzip"

    def supports(self, path: Path, payload: bytes) -> bool:
        return payload.startswith(b"\x1f\x8b") or path.suffix.lower() == ".gz"

    def detected_type(self, path: Path, payload: bytes) -> DetectedFileType:
        return DetectedFileType(name="gzip", media_type="application/gzip", extractable=True)

    def extract_members(self, task: ExtractionTask) -> list[tuple[Path, bytes]]:
        with gzip.GzipFile(fileobj=io.BytesIO(task.payload)) as archive:
            content = archive.read()
        member_name = Path(task.display_name).stem or "payload.bin"
        return [(Path(member_name), content)]


def detect_file_type(path: Path, payload: bytes) -> DetectedFileType:
    provider = detect_extraction_provider(path, payload)
    if provider is not None:
        return provider.detected_type(path, payload)
    suffix = path.suffix.lower()
    media_type = {
        ".json": "application/json",
        ".xml": "application/xml",
        ".pem": "application/x-pem-file",
        ".crt": "application/x-x509-ca-cert",
        ".cer": "application/pkix-cert",
        ".conf": "text/plain",
        ".cfg": "text/plain",
        ".ini": "text/plain",
        ".yaml": "text/yaml",
        ".yml": "text/yaml",
        ".txt": "text/plain",
        ".html": "text/html",
        ".js": "application/javascript",
        ".sh": "text/x-shellscript",
    }.get(suffix, "application/octet-stream")
    return DetectedFileType(name=suffix.lstrip(".") or "binary", media_type=media_type, extractable=False)


def detect_extraction_provider(path: Path, payload: bytes) -> ExtractionProvider | None:
    for provider in _providers():
        try:
            if provider.supports(path, payload):
                return provider
        except (OSError, zipfile.BadZipFile, tarfile.TarError, EOFError):
            continue
    return None


def extract_firmware_bundle(request: ExtractionRequest) -> ExtractionResult | None:
    payload = request.source_path.read_bytes()
    provider = detect_extraction_provider(request.source_path, payload)
    if provider is None:
        return None

    destination = request.destination_root / provider.name
    _reset_destination(destination)
    state = ExtractionState(artifacts=[], extraction_notes=[])
    queue: deque[tuple[ExtractionProvider, ExtractionTask]] = deque(
        [(provider, ExtractionTask(request.source_path.name, payload, Path(), 0))]
    )

    while queue and len(state.artifacts) < request.max_files and state.total_bytes < request.max_total_bytes:
        current_provider, task = queue.popleft()
        members = current_provider.extract_members(task)
        for relative_member, member_payload in members:
            if len(state.artifacts) >= request.max_files:
                state.truncated = True
                state.add_note("Extraction stopped after reaching the configured file-count limit.")
                break

            target_relative = task.relative_root / relative_member
            member_size = len(member_payload)
            if state.total_bytes + member_size > request.max_total_bytes:
                state.truncated = True
                state.add_note("Extraction stopped after reaching the configured byte limit.")
                break

            target = destination / target_relative
            target.parent.mkdir(parents=True, exist_ok=True)
            target.write_bytes(member_payload)
            detected = detect_file_type(Path(relative_member.name), member_payload)
            state.total_bytes += member_size
            state.artifacts.append(
                _build_artifact_record(
                    target,
                    destination,
                    current_provider.name,
                    media_type=detected.media_type,
                    extracted_from=task.extracted_from,
                    depth=task.depth,
                )
            )

            nested_provider = detect_extraction_provider(Path(relative_member.name), member_payload)
            if nested_provider is not None:
                state.nested_archives_seen += 1
                if task.depth + 1 >= request.max_depth:
                    state.depth_limit_hit = True
                    state.add_note(
                        f"Nested archive {target_relative.as_posix()} was detected but not expanded beyond depth {request.max_depth}."
                    )
                    continue
                queue.append(
                    (
                        nested_provider,
                        ExtractionTask(
                            display_name=relative_member.name,
                            payload=member_payload,
                            relative_root=target_relative.with_suffix(""),
                            depth=task.depth + 1,
                            extracted_from=target_relative.as_posix(),
                        ),
                    )
                )

    metadata = {
        "provider": provider.name,
        "detected_type": provider.detected_type(request.source_path, payload).name,
        "member_count": len(state.artifacts),
        "total_bytes": state.total_bytes,
        "max_files": request.max_files,
        "max_total_bytes": request.max_total_bytes,
        "truncated": state.truncated,
        "depth_limit_hit": state.depth_limit_hit,
        "nested_archives_seen": state.nested_archives_seen,
        "extraction_notes": sorted(state.extraction_notes or []),
    }
    return ExtractionResult(
        provider=provider.name,
        extracted_root=destination,
        artifacts=sorted(state.artifacts, key=lambda item: (item.relative_path, item.depth, item.source)),
        metadata=metadata,
    )


def _providers() -> list[ExtractionProvider]:
    return [ZipExtractionProvider(), TarExtractionProvider(), GzipExtractionProvider()]


def _safe_member_path(raw_name: str) -> Path:
    normalized = PurePosixPath(raw_name)
    parts = [part for part in normalized.parts if part not in {"", ".", ".."}]
    if not parts:
        raise ValueError("Archive member does not contain a safe relative path.")
    return Path(*parts)


def _build_artifact_record(
    path: Path,
    root: Path,
    provider: str,
    *,
    media_type: str,
    extracted_from: str | None,
    depth: int,
) -> ExtractedArtifact:
    digest = hashlib.sha256(path.read_bytes()).hexdigest()
    return ExtractedArtifact(
        relative_path=str(path.relative_to(root)).replace("\\", "/"),
        source=provider,
        kind="file",
        size_bytes=path.stat().st_size,
        sha256=digest,
        media_type=media_type,
        extracted_from=extracted_from,
        depth=depth,
        confidence=0.9,
    )


def _reset_destination(path: Path) -> None:
    if path.exists():
        shutil.rmtree(path)
    path.mkdir(parents=True, exist_ok=True)

