"""Zone gallery — file IO over data/zones/zone_{zid}/place_{pid}/.

A two-level hierarchy that replaces the flat place gallery:

    Zone   = name + description (a region, e.g. "Ground Floor")
      └─ Place = name + description + optional reference photos
                 + optional linked face ids (people associated with the place)

Layout:
    zones/
      zone_{zid}/
        meta.json                  {name, description, added_at}
        place_{pid}/
          meta.json                {name, description, face_ids:[int], added_at}
          place_1.jpg              ← optional reference photos (0..N)
          place_2.png

`face_ids` reference enrolled faces in the SEPARATE face gallery
(data/faces/face_{id}); this module only stores the ids — name/photo
resolution is done by the caller (route layer + Gemini primer).

Reference photos let Gemini visually recognise the place (and, later, let the
robot navigate to it). A place needs only a name + description; photos and
linked faces are both optional. Thread-safe via a single internal RLock.
"""

from __future__ import annotations

import io
import json
import re
import threading
import zipfile
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Iterable

from Project.Sanad.core.logger import get_logger

log = get_logger("zone_gallery")

_ZONE_DIR_RE = re.compile(r"^zone_(\d+)$")
_PLACE_DIR_RE = re.compile(r"^place_(\d+)$")
ALLOWED_EXTS = {".jpg", ".jpeg", ".png"}
SAMPLE_NAME_RE = re.compile(r"^place_(\d+)\.(jpg|jpeg|png)$", re.IGNORECASE)


@dataclass
class PlaceEntry:
    id: int
    zone_id: int
    name: str | None
    added_at: str | None
    dir: Path
    description: str | None = None
    face_ids: list[int] = field(default_factory=list)
    sample_paths: list[Path] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "zone_id": self.zone_id,
            "name": self.name,
            "description": self.description,
            "face_ids": list(self.face_ids),
            "added_at": self.added_at,
            "photos": [
                {"name": p.name, "size_bytes": p.stat().st_size}
                for p in self.sample_paths if p.exists()
            ],
        }


@dataclass
class ZoneEntry:
    id: int
    name: str | None
    added_at: str | None
    dir: Path
    description: str | None = None
    places: list[PlaceEntry] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "name": self.name,
            "description": self.description,
            "added_at": self.added_at,
            "places": [p.to_dict() for p in self.places],
        }


class ZoneGallery:
    """File-system backed zone/place gallery rooted at `root` (data/zones/)."""

    def __init__(self, root: Path | str) -> None:
        self.root = Path(root)
        self._lock = threading.RLock()

    # ── paths ────────────────────────────────────────────────

    def _ensure_root(self) -> None:
        self.root.mkdir(parents=True, exist_ok=True)

    def _zone_dir(self, zone_id: int) -> Path:
        return self.root / f"zone_{zone_id}"

    def _place_dir(self, zone_id: int, place_id: int) -> Path:
        return self.root / f"zone_{zone_id}" / f"place_{place_id}"

    def _iter_zone_dirs(self) -> Iterable[tuple[int, Path]]:
        if not self.root.exists():
            return
        for child in sorted(self.root.iterdir()):
            if not child.is_dir():
                continue
            m = _ZONE_DIR_RE.match(child.name)
            if m:
                yield int(m.group(1)), child

    def _iter_place_dirs(self, zone_dir: Path) -> Iterable[tuple[int, Path]]:
        if not zone_dir.exists():
            return
        for child in sorted(zone_dir.iterdir()):
            if not child.is_dir():
                continue
            m = _PLACE_DIR_RE.match(child.name)
            if m:
                yield int(m.group(1)), child

    def _samples_in(self, place_dir: Path) -> list[Path]:
        return [p for p in sorted(place_dir.iterdir())
                if p.is_file() and p.suffix.lower() in ALLOWED_EXTS]

    # ── meta ─────────────────────────────────────────────────

    def _zone_meta(self, zone_dir: Path) -> tuple[str | None, str | None, str | None]:
        meta_path = zone_dir / "meta.json"
        if not meta_path.exists():
            return None, None, None
        try:
            data = json.loads(meta_path.read_text(encoding="utf-8"))
        except Exception:
            return None, None, None
        return (data.get("name") or None), (data.get("description") or None), data.get("added_at")

    def _place_meta(self, place_dir: Path) -> tuple[str | None, str | None, list[int], str | None]:
        meta_path = place_dir / "meta.json"
        if not meta_path.exists():
            return None, None, [], None
        try:
            data = json.loads(meta_path.read_text(encoding="utf-8"))
        except Exception:
            return None, None, [], None
        raw_ids = data.get("face_ids") or []
        face_ids: list[int] = []
        for x in raw_ids:
            try:
                face_ids.append(int(x))
            except (TypeError, ValueError):
                continue
        return (data.get("name") or None), (data.get("description") or None), face_ids, data.get("added_at")

    def _write_zone_meta(self, zone_dir: Path, name, description, added_at=None) -> None:
        meta: dict = {}
        if name:
            meta["name"] = name
        if description:
            meta["description"] = description
        meta["added_at"] = added_at or datetime.now().isoformat(timespec="seconds")
        (zone_dir / "meta.json").write_text(
            json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")

    def _write_place_meta(self, place_dir: Path, name, description,
                          face_ids=None, added_at=None) -> None:
        meta: dict = {}
        if name:
            meta["name"] = name
        if description:
            meta["description"] = description
        meta["face_ids"] = [int(x) for x in (face_ids or [])]
        meta["added_at"] = added_at or datetime.now().isoformat(timespec="seconds")
        (place_dir / "meta.json").write_text(
            json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")

    # ── read ─────────────────────────────────────────────────

    def _build_place(self, zone_id: int, place_id: int, place_dir: Path) -> PlaceEntry:
        name, desc, face_ids, added = self._place_meta(place_dir)
        return PlaceEntry(
            id=place_id, zone_id=zone_id, name=name, description=desc,
            face_ids=face_ids, added_at=added, dir=place_dir,
            sample_paths=self._samples_in(place_dir),
        )

    def _build_zone(self, zone_id: int, zone_dir: Path) -> ZoneEntry:
        name, desc, added = self._zone_meta(zone_dir)
        places = [self._build_place(zone_id, pid, pdir)
                  for pid, pdir in self._iter_place_dirs(zone_dir)]
        return ZoneEntry(id=zone_id, name=name, description=desc,
                         added_at=added, dir=zone_dir, places=places)

    def list_zones(self) -> list[ZoneEntry]:
        with self._lock:
            return [self._build_zone(zid, zdir) for zid, zdir in self._iter_zone_dirs()]

    def get_zone(self, zone_id: int) -> ZoneEntry | None:
        with self._lock:
            zd = self._zone_dir(zone_id)
            return self._build_zone(zone_id, zd) if zd.is_dir() else None

    def get_place(self, zone_id: int, place_id: int) -> PlaceEntry | None:
        with self._lock:
            pd = self._place_dir(zone_id, place_id)
            return self._build_place(zone_id, place_id, pd) if pd.is_dir() else None

    def get_photo(self, zone_id: int, place_id: int, photo_name: str) -> Path | None:
        with self._lock:
            pd = self._place_dir(zone_id, place_id)
            if not pd.is_dir():
                return None
            p = pd / photo_name
            try:
                p.resolve().relative_to(pd.resolve())
            except ValueError:
                return None
            if not p.exists() or p.suffix.lower() not in ALLOWED_EXTS:
                return None
            return p

    # ── ids ──────────────────────────────────────────────────

    def next_zone_id(self) -> int:
        with self._lock:
            return max((zid for zid, _ in self._iter_zone_dirs()), default=0) + 1

    def next_place_id(self, zone_id: int) -> int:
        with self._lock:
            zd = self._zone_dir(zone_id)
            return max((pid for pid, _ in self._iter_place_dirs(zd)), default=0) + 1

    def _next_sample_name(self, place_dir: Path, ext: str) -> str:
        max_n = 0
        for p in self._samples_in(place_dir):
            m = SAMPLE_NAME_RE.match(p.name)
            if m:
                max_n = max(max_n, int(m.group(1)))
        return f"place_{max_n + 1}{ext.lower()}"

    @staticmethod
    def _detect_ext(data: bytes) -> str:
        if len(data) >= 8 and data[:8] == b"\x89PNG\r\n\x1a\n":
            return ".png"
        return ".jpg"

    # ── zone write ───────────────────────────────────────────

    def create_zone(self, name: str | None = None,
                    description: str | None = None) -> ZoneEntry:
        with self._lock:
            self._ensure_root()
            zid = self.next_zone_id()
            zd = self._zone_dir(zid)
            zd.mkdir(parents=True, exist_ok=False)
            self._write_zone_meta(zd, (name or "").strip() or None,
                                  (description or "").strip() or None)
            log.info("Created zone_%d (name=%s)", zid, name or "(unnamed)")
            return self._build_zone(zid, zd)

    def rename_zone(self, zone_id: int, name: str | None) -> None:
        with self._lock:
            zd = self._zone_dir(zone_id)
            if not zd.is_dir():
                raise FileNotFoundError(f"zone_{zone_id} not found")
            _, desc, added = self._zone_meta(zd)
            self._write_zone_meta(zd, (name or "").strip() or None, desc, added_at=added)
            log.info("Renamed zone_%d → %s", zone_id, name or "(unnamed)")

    def describe_zone(self, zone_id: int, description: str | None) -> None:
        with self._lock:
            zd = self._zone_dir(zone_id)
            if not zd.is_dir():
                raise FileNotFoundError(f"zone_{zone_id} not found")
            name, _, added = self._zone_meta(zd)
            self._write_zone_meta(zd, name, (description or "").strip() or None, added_at=added)
            log.info("Described zone_%d", zone_id)

    def delete_zone(self, zone_id: int) -> None:
        import shutil
        with self._lock:
            zd = self._zone_dir(zone_id)
            if not zd.is_dir():
                raise FileNotFoundError(f"zone_{zone_id} not found")
            shutil.rmtree(zd)
            log.info("Deleted zone_%d (and its places)", zone_id)

    # ── place write ──────────────────────────────────────────

    def create_place(self, zone_id: int, name: str | None = None,
                     description: str | None = None,
                     face_ids: list[int] | None = None,
                     image_bytes_list: list[bytes] | None = None) -> PlaceEntry:
        with self._lock:
            zd = self._zone_dir(zone_id)
            if not zd.is_dir():
                raise FileNotFoundError(f"zone_{zone_id} not found")
            pid = self.next_place_id(zone_id)
            pd = self._place_dir(zone_id, pid)
            pd.mkdir(parents=True, exist_ok=False)
            for idx, data in enumerate(image_bytes_list or [], start=1):
                (pd / f"place_{idx}{self._detect_ext(data)}").write_bytes(data)
            self._write_place_meta(pd, (name or "").strip() or None,
                                   (description or "").strip() or None, face_ids or [])
            log.info("Created zone_%d/place_%d (name=%s, photos=%d, faces=%d)",
                     zone_id, pid, name or "(unnamed)",
                     len(image_bytes_list or []), len(face_ids or []))
            return self._build_place(zone_id, pid, pd)

    def rename_place(self, zone_id: int, place_id: int, name: str | None) -> None:
        with self._lock:
            pd = self._place_dir(zone_id, place_id)
            if not pd.is_dir():
                raise FileNotFoundError(f"zone_{zone_id}/place_{place_id} not found")
            _, desc, fids, added = self._place_meta(pd)
            self._write_place_meta(pd, (name or "").strip() or None, desc, fids, added_at=added)
            log.info("Renamed zone_%d/place_%d → %s", zone_id, place_id, name or "(unnamed)")

    def describe_place(self, zone_id: int, place_id: int, description: str | None) -> None:
        with self._lock:
            pd = self._place_dir(zone_id, place_id)
            if not pd.is_dir():
                raise FileNotFoundError(f"zone_{zone_id}/place_{place_id} not found")
            name, _, fids, added = self._place_meta(pd)
            self._write_place_meta(pd, name, (description or "").strip() or None, fids, added_at=added)
            log.info("Described zone_%d/place_%d", zone_id, place_id)

    def set_place_faces(self, zone_id: int, place_id: int, face_ids: list[int]) -> None:
        """Replace the set of linked face ids for a place."""
        with self._lock:
            pd = self._place_dir(zone_id, place_id)
            if not pd.is_dir():
                raise FileNotFoundError(f"zone_{zone_id}/place_{place_id} not found")
            name, desc, _, added = self._place_meta(pd)
            clean = []
            for x in (face_ids or []):
                try:
                    clean.append(int(x))
                except (TypeError, ValueError):
                    continue
            self._write_place_meta(pd, name, desc, clean, added_at=added)
            log.info("Set zone_%d/place_%d faces → %s", zone_id, place_id, clean)

    def add_photo(self, zone_id: int, place_id: int, image_bytes: bytes) -> str:
        with self._lock:
            pd = self._place_dir(zone_id, place_id)
            if not pd.is_dir():
                raise FileNotFoundError(f"zone_{zone_id}/place_{place_id} not found")
            fname = self._next_sample_name(pd, self._detect_ext(image_bytes))
            (pd / fname).write_bytes(image_bytes)
            log.info("Added %s to zone_%d/place_%d", fname, zone_id, place_id)
            return fname

    def delete_photo(self, zone_id: int, place_id: int, photo_name: str) -> None:
        with self._lock:
            pd = self._place_dir(zone_id, place_id)
            if not pd.is_dir():
                raise FileNotFoundError(f"zone_{zone_id}/place_{place_id} not found")
            target = self.get_photo(zone_id, place_id, photo_name)
            if target is None:
                raise FileNotFoundError(f"photo {photo_name} not found")
            target.unlink()
            log.info("Deleted %s from zone_%d/place_%d", photo_name, zone_id, place_id)

    def delete_place(self, zone_id: int, place_id: int) -> None:
        import shutil
        with self._lock:
            pd = self._place_dir(zone_id, place_id)
            if not pd.is_dir():
                raise FileNotFoundError(f"zone_{zone_id}/place_{place_id} not found")
            shutil.rmtree(pd)
            log.info("Deleted zone_%d/place_%d", zone_id, place_id)

    def zip_place(self, zone_id: int, place_id: int) -> bytes:
        with self._lock:
            pd = self._place_dir(zone_id, place_id)
            if not pd.is_dir():
                raise FileNotFoundError(f"zone_{zone_id}/place_{place_id} not found")
            buf = io.BytesIO()
            with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as zf:
                for p in sorted(pd.iterdir()):
                    if p.is_file():
                        zf.write(p, arcname=f"zone_{zone_id}_place_{place_id}/{p.name}")
            return buf.getvalue()

    # ── primer support (used by gemini/script.py) ───────────

    def load_for_primer(self, max_samples_per_place: int = 3,
                        resize_long_side: int = 256
                        ) -> list[tuple[ZoneEntry, list[tuple[PlaceEntry, list[bytes]]]]]:
        """Return [(ZoneEntry, [(PlaceEntry, [jpeg_bytes,…]), …]), …].

        Photo lists may be empty (name+description-only place). Photos are
        resized to longest-side <= resize_long_side, re-encoded JPEG q=85.
        """
        out: list[tuple[ZoneEntry, list[tuple[PlaceEntry, list[bytes]]]]] = []
        for zone in self.list_zones():
            place_jpegs: list[tuple[PlaceEntry, list[bytes]]] = []
            for place in zone.places:
                jpegs: list[bytes] = []
                for p in place.sample_paths[:max_samples_per_place]:
                    try:
                        raw = p.read_bytes()
                    except OSError:
                        continue
                    jpegs.append(self._resize_for_primer(raw, resize_long_side) or raw)
                place_jpegs.append((place, jpegs))
            out.append((zone, place_jpegs))
        return out

    @staticmethod
    def _resize_for_primer(raw: bytes, long_side: int) -> bytes | None:
        try:
            from PIL import Image  # type: ignore
        except Exception:
            return None
        try:
            img = Image.open(io.BytesIO(raw))
            img.load()
            if img.mode not in ("RGB", "L"):
                img = img.convert("RGB")
            w, h = img.size
            scale = long_side / max(w, h) if max(w, h) > long_side else 1.0
            if scale < 1.0:
                img = img.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.LANCZOS)
            buf = io.BytesIO()
            img.save(buf, format="JPEG", quality=85, optimize=True)
            return buf.getvalue()
        except Exception:
            return None