"""Worker pool — per-user orchestrators with persistent Zo workspaces.

Two storage modes:

* **zo** — User's data lives persistently on Zo at
  ``ZO_WORKSPACES_DIR/{user_id}/``.  On hydration the workspace is created
  (or reused if it already exists) and seeded.  On evict the orchestrator shuts
  down but the workspace is **not deleted**.

* **byos** — User's canonical data lives in their own cloud storage (Google
  Drive, Dropbox).  On hydration data is downloaded to a temp dir.  On evict
  changed files are uploaded back, then the temp dir is deleted.

Every user (including the owner) is a registered user with a real ID — there
is no "default" user special case.  The localhost / desktop case is handled
by the ``AuthMiddleware`` resolving to the owner's user ID.

Phase 1 — single Zo, warm pool of up to MAX_WARM_SLOTS concurrent users.
"""

from __future__ import annotations

import asyncio
import logging
import os
import shutil
import tempfile
import time
from dataclasses import dataclass, field
from datetime import UTC
from pathlib import Path
from typing import TYPE_CHECKING, Any

from vastian.config import Settings, make_user_settings

if TYPE_CHECKING:
    from vastian.bridges.protocols import StorageProvider
    from vastian.orchestrator.orchestrator import Orchestrator

logger = logging.getLogger(__name__)

# ── Constants ────────────────────────────────────────────────────

MAX_WARM_SLOTS = 5  # Max simultaneous hydrated users on Zo
IDLE_EVICT_SECONDS = 600  # 10 min idle → evict orchestrator (workspace stays)

# Persistent workspace root on Zo.  Each user gets a sub-directory.
ZO_WORKSPACES_DIR = Path(os.environ.get("VASTIAN_WORKSPACES_DIR", "/home/workspace/vastian-users"))

# Data root on Zo (for hydration cache, analytics, templates)
ZO_DATA_ROOT = Path(os.environ.get("VASTIAN_DATA_ROOT", "/home/workspace/vastian-data"))

# Template workspace (copied for new users)
ZO_WORKSPACE_TEMPLATE = ZO_DATA_ROOT / "workspace-template"

# Where agents/ YAMLs live (in the app repo, NOT per-user)
_AGENTS_DIR: Path | None = None


def _get_agents_dir() -> Path:
    """Resolve the agents/ directory (app-level, shared across users)."""
    global _AGENTS_DIR  # noqa: PLW0603
    if _AGENTS_DIR is None:
        here = Path(__file__).resolve()
        for p in here.parents:
            if (p / "agents").is_dir() and (p / "pyproject.toml").exists():
                _AGENTS_DIR = p / "agents"
                break
        if _AGENTS_DIR is None:
            _AGENTS_DIR = here.parents[3] / "agents"
    return _AGENTS_DIR


# ── UserSlot ─────────────────────────────────────────────────────


@dataclass
class UserSlot:
    """One user's processing context on Zo."""

    user_id: str
    user_name: str = "User"
    user_tier: str = "open"
    storage_mode: str = "zo"  # "zo" (persistent) or "byos" (ephemeral)
    data_root: Path | None = None
    settings: Settings | None = None
    orchestrator: Any = None  # Orchestrator — typed as Any to avoid circular import
    storage_provider: Any = None  # StorageProvider (only for byos mode)
    last_active: float = field(default_factory=time.monotonic)
    hydration_ts: float = field(
        default_factory=time.time
    )  # wall-clock when slot was created/hydrated
    _initialized: bool = False

    @property
    def is_zo_resident(self) -> bool:
        return self.storage_mode == "zo"

    @property
    def idle_seconds(self) -> float:
        return time.monotonic() - self.last_active

    def touch(self) -> None:
        """Mark this slot as active (resets idle timer)."""
        self.last_active = time.monotonic()


# ── WorkerPool ───────────────────────────────────────────────────


class WorkerPool:
    """Manages per-user orchestrators with Zo-persistent or BYOS-ephemeral storage.

    Usage::

        pool = WorkerPool()

        # Zo-resident user (persistent workspace):
        orch = await pool.get_orchestrator("abc123", user_name="Alice", user_tier="admin")

        # BYOS user (ephemeral, synced from cloud):
        orch = await pool.get_orchestrator("def456", storage_provider=gdrive, storage_mode="byos")
    """

    def __init__(self, max_slots: int = MAX_WARM_SLOTS) -> None:
        self.max_slots = max_slots
        self._slots: dict[str, UserSlot] = {}
        self._lock = asyncio.Lock()

    @property
    def active_users(self) -> list[str]:
        """Return list of currently hydrated user IDs."""
        return list(self._slots.keys())

    @property
    def slot_count(self) -> int:
        """Number of warm slots."""
        return len(self._slots)

    async def get_orchestrator(
        self,
        user_id: str,
        *,
        user_name: str = "User",
        user_tier: str = "open",
        storage_mode: str = "zo",
        storage_provider: StorageProvider | None = None,
    ) -> Orchestrator:
        """Get or create an orchestrator for the given user.

        All users go through the same path — persistent Zo workspace or BYOS.
        """
        async with self._lock:
            # Existing warm slot?
            if user_id in self._slots:
                slot = self._slots[user_id]
                slot.touch()
                if slot.orchestrator is not None:
                    return slot.orchestrator

            # Need a new slot — evict if at capacity
            if len(self._slots) >= self.max_slots:
                await self._evict_lru()

            # Hydrate
            slot = await self._hydrate(
                user_id,
                user_name=user_name,
                user_tier=user_tier,
                storage_mode=storage_mode,
                storage_provider=storage_provider,
            )
            self._slots[user_id] = slot
            return slot.orchestrator

    def get_slot(self, user_id: str) -> UserSlot | None:
        """Look up a user slot (for inbox routing, etc.)."""
        return self._slots.get(user_id)

    async def flush_user(self, user_id: str) -> None:
        """Flush a user's changes back to BYOS (if applicable) and evict."""
        async with self._lock:
            slot = self._slots.pop(user_id, None)
            if slot:
                await self._flush(slot)
                await self._evict(slot)

    async def flush_all(self) -> None:
        """Flush all users (called on shutdown)."""
        async with self._lock:
            for user_id in list(self._slots.keys()):
                slot = self._slots.pop(user_id)
                await self._flush(slot)
                await self._evict(slot)

    async def evict_idle(self) -> list[str]:
        """Evict users that have been idle too long.  Returns evicted user IDs."""
        evicted: list[str] = []
        async with self._lock:
            for user_id in list(self._slots.keys()):
                slot = self._slots[user_id]
                if slot.idle_seconds > IDLE_EVICT_SECONDS:
                    self._slots.pop(user_id)
                    await self._flush(slot)
                    await self._evict(slot)
                    evicted.append(user_id)
        return evicted

    def status(self) -> dict[str, Any]:
        """Return pool status for the /api/worker/status endpoint."""
        return {
            "max_slots": self.max_slots,
            "active_slots": self.slot_count,
            "active_users": self.active_users,
            "workspaces_dir": str(ZO_WORKSPACES_DIR),
            "data_root": str(ZO_DATA_ROOT),
            "template_available": ZO_WORKSPACE_TEMPLATE.is_dir(),
            "slots": {
                uid: {
                    "user_name": s.user_name,
                    "user_tier": s.user_tier,
                    "storage_mode": s.storage_mode,
                    "idle_seconds": round(s.idle_seconds, 1),
                    "initialized": s._initialized,
                }
                for uid, s in self._slots.items()
            },
        }

    # ── Internal lifecycle ────────────────────────────────────────

    async def _hydrate(
        self,
        user_id: str,
        *,
        user_name: str,
        user_tier: str,
        storage_mode: str,
        storage_provider: StorageProvider | None,
    ) -> UserSlot:
        """Create and hydrate a user slot.

        Zo-resident users get a persistent workspace at ZO_WORKSPACES_DIR/{user_id}/.
        BYOS users get a temp dir that is deleted on evict.
        """
        from vastian.orchestrator.orchestrator import Orchestrator
        from vastian.seed import hydrate_user_data

        # 1. Resolve data root
        if storage_mode == "zo":
            data_root = ZO_WORKSPACES_DIR / user_id
            is_new = not data_root.exists()
            data_root.mkdir(parents=True, exist_ok=True)
            # If brand-new workspace, seed from template
            if is_new and ZO_WORKSPACE_TEMPLATE.is_dir():
                logger.info("Seeding new Zo workspace from template for %s", user_id)
                self._seed_from_template(data_root, user_id)
            else:
                # Ensure standard subdirectories exist even for existing workspaces
                for sub in ("docs", "state", "config/local", "saved_prompts/inbox", "uploads"):
                    (data_root / sub).mkdir(parents=True, exist_ok=True)
            logger.info("Hydrating Zo-resident user %s at %s", user_id, data_root)
        else:
            data_root = Path(
                tempfile.mkdtemp(
                    prefix=f"vastian-user-{user_id}-",
                    dir=str(ZO_DATA_ROOT / "hydration-cache")
                    if (ZO_DATA_ROOT / "hydration-cache").is_dir()
                    else None,
                )
            )
            # Seed from template for BYOS too
            if ZO_WORKSPACE_TEMPLATE.is_dir():
                self._seed_from_template(data_root, user_id)
            logger.info("Hydrating BYOS user %s at %s", user_id, data_root)
            # Download from BYOS (overwrites template with real data)
            if storage_provider is not None:
                await self._download_from_byos(storage_provider, data_root)

        # 2. Approach C: config seed, doc self-heal, state fixups, migrations
        agents_dir = _get_agents_dir()
        result = hydrate_user_data(agents_dir, data_root)
        if result.get("docs_healed"):
            logger.info("User %s: healed %d docs", user_id, len(result["docs_healed"]))
        if result.get("migrations_applied"):
            logger.info(
                "User %s: applied %d migrations", user_id, len(result["migrations_applied"])
            )

        # 3. Per-user Settings + Orchestrator
        settings = make_user_settings(
            data_root,
            user_name=user_name,
            user_tier=user_tier,
        )
        orch = Orchestrator(settings=settings, user_id=user_id)
        await orch.initialize(session_only=True)

        slot = UserSlot(
            user_id=user_id,
            user_name=user_name,
            user_tier=user_tier,
            storage_mode=storage_mode,
            data_root=data_root,
            settings=settings,
            orchestrator=orch,
            storage_provider=storage_provider,
            _initialized=True,
        )
        return slot

    async def _flush(self, slot: UserSlot) -> None:
        """Upload changed files back to the user's BYOS provider (BYOS mode only).

        Uses delta-sync: only uploads files whose mtime is newer than the slot's
        hydration timestamp, so we don't re-upload unchanged seed data.
        """
        if slot.storage_mode != "byos" or slot.storage_provider is None or slot.data_root is None:
            return
        try:
            await self._upload_to_byos(
                slot.storage_provider, slot.data_root, since=slot.hydration_ts
            )
            logger.info(
                "Flushed user %s to BYOS (delta since %.0f)", slot.user_id, slot.hydration_ts
            )
        except Exception:
            logger.error("Failed to flush user %s to BYOS", slot.user_id, exc_info=True)

    async def _evict(self, slot: UserSlot) -> None:
        """Clean up an evicted user slot.

        Zo-resident: orchestrator shutdown only, workspace persists.
        BYOS: orchestrator shutdown + delete temp dir.
        """
        # Shut down orchestrator
        if slot.orchestrator is not None:
            try:
                await slot.orchestrator.shutdown()
            except Exception:
                logger.warning("Orchestrator shutdown failed for %s", slot.user_id, exc_info=True)

        # Only delete ephemeral dirs (BYOS mode)
        if slot.storage_mode == "byos" and slot.data_root and slot.data_root.exists():
            try:
                shutil.rmtree(slot.data_root, ignore_errors=True)
                logger.info("Evicted BYOS user %s, cleaned %s", slot.user_id, slot.data_root)
            except Exception:
                logger.warning("Cleanup failed for %s", slot.user_id, exc_info=True)
        elif slot.is_zo_resident:
            logger.info(
                "Evicted Zo-resident user %s (workspace preserved at %s)",
                slot.user_id,
                slot.data_root,
            )

    async def _evict_lru(self) -> None:
        """Evict the least recently used slot."""
        if not self._slots:
            return
        lru_id = min(self._slots, key=lambda uid: self._slots[uid].last_active)
        slot = self._slots.pop(lru_id)
        await self._flush(slot)
        await self._evict(slot)
        logger.info("LRU-evicted user %s (idle %.0fs)", lru_id, slot.idle_seconds)

    # ── Workspace template seeding ─────────────────────────────────

    @staticmethod
    def _seed_from_template(data_root: Path, user_id: str) -> None:
        """Copy the workspace template into a new user's data root.

        Replaces ``__USER_ID__`` and ``__CREATED_AT__`` placeholders in
        template files with actual values.
        """
        import json
        from datetime import datetime

        template = ZO_WORKSPACE_TEMPLATE
        if not template.is_dir():
            return

        # Walk and copy directory structure + files
        for src_path in template.rglob("*"):
            rel = src_path.relative_to(template)
            dest = data_root / rel
            if src_path.is_dir():
                dest.mkdir(parents=True, exist_ok=True)
            else:
                dest.parent.mkdir(parents=True, exist_ok=True)
                if not dest.exists():
                    content = src_path.read_text(encoding="utf-8", errors="replace")
                    content = content.replace("__USER_ID__", user_id)
                    content = content.replace(
                        "__CREATED_AT__",
                        datetime.now(UTC).isoformat(),
                    )
                    dest.write_text(content, encoding="utf-8")

        # Ensure identity.json exists with real user_id
        identity_path = data_root / "identity.json"
        if not identity_path.exists():
            identity_path.write_text(
                json.dumps(
                    {
                        "schema_version": 1,
                        "user_id": user_id,
                        "created_at": datetime.now(UTC).isoformat(),
                        "devices": [],
                        "storage_mode": "zo",
                    },
                    indent=2,
                ),
                encoding="utf-8",
            )

        logger.info("Seeded workspace for %s from template (%s)", user_id, template)

    # ── BYOS transfer helpers ─────────────────────────────────────

    async def _download_from_byos(
        self,
        provider: StorageProvider,
        data_root: Path,
    ) -> None:
        """Download user data from BYOS to local dir.

        Only downloads docs/, state/, config/, saved_prompts/ — not data/ (SQLite, ChromaDB).
        """
        sync_dirs = ["docs", "state", "config", "saved_prompts"]

        for subdir in sync_dirs:
            try:
                remote_prefix = f"live/{subdir}"
                files = provider.list_files(remote_prefix)
                for remote_path in files:
                    rel = remote_path.removeprefix("live/")
                    local = data_root / rel
                    local.parent.mkdir(parents=True, exist_ok=True)
                    content = provider.download(remote_path)
                    if content is not None:
                        local.write_bytes(content)
            except Exception:
                logger.warning(
                    "BYOS download failed for %s (user data may be partial)",
                    subdir,
                    exc_info=True,
                )

    async def _upload_to_byos(
        self,
        provider: StorageProvider,
        data_root: Path,
        *,
        since: float = 0.0,
    ) -> None:
        """Upload changed files from local dir back to BYOS.

        Args:
            since: Only upload files modified after this UNIX timestamp (delta sync).
                   0 means upload everything.
        """
        sync_dirs = ["docs", "state", "config", "saved_prompts"]
        uploaded = 0

        for subdir in sync_dirs:
            local_dir = data_root / subdir
            if not local_dir.exists():
                continue
            for local_path in local_dir.rglob("*"):
                if local_path.is_dir():
                    continue
                if local_path.name.startswith("."):
                    continue
                # Delta-sync: skip files not modified since hydration
                if since and local_path.stat().st_mtime < since:
                    continue
                rel = local_path.relative_to(data_root)
                remote = f"live/{rel.as_posix()}"
                try:
                    provider.upload(remote, local_path.read_bytes())
                    uploaded += 1
                except Exception:
                    logger.warning("BYOS upload failed: %s", remote, exc_info=True)

        logger.debug("BYOS delta-sync: uploaded %d files (since=%.0f)", uploaded, since)
