"""JWT authentication middleware and user registry for cloud multi-user.

Users register with email + password and get a JWT.  The JWT is validated
on every API request.  Localhost / desktop requests resolve to the **owner**
(first admin-tier user) rather than a hardcoded "default" user.

The registry lives at ``VASTIAN_DATA_ROOT/vastian-registry/users.json``
on the Zo host (NOT in any user's BYOS).

BYOS tokens (OAuth tokens for Google Drive, Dropbox, etc.) are stored
encrypted alongside the user record.  Phase 1 uses Fernet symmetric
encryption with a server-side key.
"""

from __future__ import annotations

import hashlib
import hmac
import json
import logging
import os
import secrets
import shutil
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.responses import JSONResponse, RedirectResponse, Response

logger = logging.getLogger(__name__)

# ── JWT (HS256, no external dep) ─────────────────────────────────

_JWT_ALG = "HS256"
_JWT_EXPIRY = 86400  # 24 hours


def _get_jwt_secret() -> str:
    """Return the JWT signing secret.  Reads from env or generates one."""
    secret = os.environ.get("VASTIAN_JWT_SECRET", "")
    if not secret:
        secret_path = _registry_dir() / ".jwt_secret"
        if secret_path.exists():
            secret = secret_path.read_text(encoding="utf-8").strip()
        else:
            secret = secrets.token_hex(32)
            secret_path.parent.mkdir(parents=True, exist_ok=True)
            secret_path.write_text(secret, encoding="utf-8")
    return secret


def _b64url_encode(data: bytes) -> str:
    import base64

    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def _b64url_decode(s: str) -> bytes:
    import base64

    padding = 4 - len(s) % 4
    return base64.urlsafe_b64decode(s + "=" * padding)


def create_jwt(user_id: str, extra: dict[str, Any] | None = None) -> str:
    """Create a signed JWT token for a user."""
    header = _b64url_encode(json.dumps({"alg": _JWT_ALG, "typ": "JWT"}).encode())
    now = int(time.time())
    payload_data: dict[str, Any] = {
        "sub": user_id,
        "iat": now,
        "exp": now + _JWT_EXPIRY,
    }
    if extra:
        payload_data.update(extra)
    payload = _b64url_encode(json.dumps(payload_data).encode())
    sig_input = f"{header}.{payload}".encode()
    signature = hmac.new(_get_jwt_secret().encode(), sig_input, hashlib.sha256).digest()
    sig = _b64url_encode(signature)
    return f"{header}.{payload}.{sig}"


def verify_jwt(token: str) -> dict[str, Any] | None:
    """Verify a JWT and return the payload, or None if invalid/expired."""
    try:
        parts = token.split(".")
        if len(parts) != 3:
            return None
        header_b64, payload_b64, sig_b64 = parts
        sig_input = f"{header_b64}.{payload_b64}".encode()
        expected = hmac.new(_get_jwt_secret().encode(), sig_input, hashlib.sha256).digest()
        actual = _b64url_decode(sig_b64)
        if not hmac.compare_digest(expected, actual):
            return None
        payload = json.loads(_b64url_decode(payload_b64))
        if payload.get("exp", 0) < time.time():
            return None
        return payload
    except Exception:
        return None


# ── User Registry ────────────────────────────────────────────────


def _registry_dir() -> Path:
    """Directory for the server-side user registry."""
    data_root = os.environ.get("VASTIAN_DATA_ROOT", "")
    if data_root:
        return Path(data_root) / "vastian-registry"
    # Fallback to platform default
    from vastian.storage.data_root import get_data_root

    return get_data_root() / "vastian-registry"


def _users_path() -> Path:
    return _registry_dir() / "users.json"


def _load_users() -> dict[str, dict[str, Any]]:
    """Load all users from the registry.  Returns {user_id: {...}}."""
    path = _users_path()
    if not path.exists():
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
        return data if isinstance(data, dict) else {}
    except Exception:
        logger.warning("Failed to load user registry", exc_info=True)
        return {}


def _save_users(users: dict[str, dict[str, Any]]) -> None:
    path = _users_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(users, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )


def _hash_password(password: str, salt: str | None = None) -> tuple[str, str]:
    """Hash a password with PBKDF2.  Returns (hash_hex, salt_hex)."""
    if salt is None:
        salt = secrets.token_hex(16)
    hashed = hashlib.pbkdf2_hmac("sha256", password.encode(), salt.encode(), 100_000)
    return hashed.hex(), salt


def _normalize_passphrase(phrase: str) -> str:
    """Normalize 12-word passphrase: lowercase, collapse whitespace, strip."""
    return " ".join(phrase.lower().split())


def _hash_passphrase(phrase: str, salt: str | None = None) -> tuple[str, str]:
    """Hash a passphrase with PBKDF2.  Returns (hash_hex, salt_hex)."""
    if salt is None:
        salt = secrets.token_hex(16)
    normalized = _normalize_passphrase(phrase)
    hashed = hashlib.pbkdf2_hmac(
        "sha256", normalized.encode(), salt.encode(), 100_000
    )
    return hashed.hex(), salt


def _derive_user_id_from_passphrase(phrase: str) -> str:
    """Derive deterministic 16-char hex user_id from passphrase."""
    normalized = _normalize_passphrase(phrase)
    digest = hashlib.sha256(normalized.encode()).hexdigest()
    return digest[:16]


@dataclass
class UserRecord:
    user_id: str
    email: str
    name: str
    tier: str
    storage_mode: str  # "zo" or "byos"
    byos_provider: str  # "gdrive" | "dropbox" | "box" | "tresorit" | ""
    byos_token_encrypted: str
    password_hash: str
    password_salt: str
    created_at: str


def register_user(
    email: str,
    password: str,
    name: str = "",
    tier: str = "open",
    storage_mode: str = "zo",
    byos_provider: str = "",
) -> tuple[str | None, str]:
    """Register a new user.  Returns (user_id, error_message).

    Returns (None, error) if email already exists.
    """
    users = _load_users()

    # Check for duplicate email
    for _uid, u in users.items():
        if u.get("email", "").lower() == email.lower():
            return None, "Email already registered"

    user_id = secrets.token_hex(8)  # 16-char hex ID
    pw_hash, pw_salt = _hash_password(password)

    users[user_id] = {
        "email": email.lower(),
        "name": name or email.split("@")[0],
        "tier": tier,
        "storage_mode": storage_mode,
        "byos_provider": byos_provider,
        "byos_token_encrypted": "",
        "auth_method": "email",
        "password_hash": pw_hash,
        "password_salt": pw_salt,
        "passphrase_hash": "",
        "passphrase_salt": "",
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }
    _save_users(users)

    # Create persistent workspace for Zo-resident users
    if storage_mode == "zo":
        from vastian.server.worker_pool import ZO_WORKSPACES_DIR

        workspace = ZO_WORKSPACES_DIR / user_id
        workspace.mkdir(parents=True, exist_ok=True)
        logger.info("Created Zo workspace for user %s at %s", user_id, workspace)

    return user_id, ""


def authenticate_user(email: str, password: str) -> tuple[str | None, dict[str, Any]]:
    """Authenticate a user by email + password.

    Returns (user_id, user_dict) on success, (None, {}) on failure.
    """
    users = _load_users()
    for uid, u in users.items():
        if u.get("email", "").lower() == email.lower():
            pw_hash, _ = _hash_password(password, u.get("password_salt", ""))
            if hmac.compare_digest(pw_hash, u.get("password_hash", "")):
                return uid, u
            return None, {}
    return None, {}


def get_user(user_id: str) -> dict[str, Any] | None:
    """Look up a user by ID."""
    users = _load_users()
    return users.get(user_id)


def authenticate_by_passphrase(phrase: str) -> tuple[str | None, dict[str, Any]]:
    """Authenticate a user by 12-word passphrase.

    Returns (user_id, user_dict) on success, (None, {}) on failure.
    Requires exactly 12 words (normalized).
    """
    normalized = _normalize_passphrase(phrase)
    words = normalized.split()
    if len(words) != 12:
        return None, {}

    user_id = _derive_user_id_from_passphrase(phrase)
    users = _load_users()
    u = users.get(user_id)
    if u is None:
        return None, {}

    auth_method = u.get("auth_method", "email")
    if auth_method != "passphrase":
        return None, {}

    pw_hash, _ = _hash_passphrase(phrase, u.get("passphrase_salt", ""))
    if hmac.compare_digest(pw_hash, u.get("passphrase_hash", "")):
        return user_id, u
    return None, {}


def bootstrap_owner_passphrase(
    phrase: str,
    name: str = "",
) -> tuple[str | None, str]:
    """Bootstrap the system by creating the first owner user from passphrase.

    Returns (user_id, error_message).
    Requires exactly 12 words.
    """
    normalized = _normalize_passphrase(phrase)
    words = normalized.split()
    if len(words) != 12:
        return None, "Passphrase must be exactly 12 words"

    users = _load_users()
    if users:
        for uid, u in users.items():
            if u.get("tier") == "admin":
                return uid, f"Owner already exists ({u.get('name', uid)})"

    user_id = _derive_user_id_from_passphrase(phrase)
    if user_id in users:
        return None, "User already exists with this passphrase"

    pw_hash, pw_salt = _hash_passphrase(phrase)

    users[user_id] = {
        "email": "",
        "name": name or "Owner",
        "tier": "admin",
        "storage_mode": "zo",
        "byos_provider": "",
        "byos_token_encrypted": "",
        "auth_method": "passphrase",
        "passphrase_hash": pw_hash,
        "passphrase_salt": pw_salt,
        "password_hash": "",
        "password_salt": "",
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }
    _save_users(users)

    from vastian.server.worker_pool import ZO_WORKSPACES_DIR

    workspace = ZO_WORKSPACES_DIR / user_id
    workspace.mkdir(parents=True, exist_ok=True)
    logger.info("Bootstrapped owner from passphrase: %s", user_id)
    return user_id, ""


def get_owner_user() -> tuple[str | None, dict[str, Any] | None]:
    """Find the owner (first admin-tier user).  Returns (user_id, user_dict) or (None, None)."""
    users = _load_users()
    for uid, u in users.items():
        if u.get("tier") == "admin":
            return uid, u
    # No admin found — return first user if any
    for uid, u in users.items():
        return uid, u
    return None, None


def bootstrap_owner(
    email: str,
    password: str,
    name: str = "",
    *,
    migrate_from: Path | None = None,
) -> tuple[str | None, str]:
    """Bootstrap the system by creating the first owner user.

    Returns (user_id, error_message).

    If ``migrate_from`` is given (e.g. an existing data root like
    ``/root/.local/share/vastian``), existing data is copied into the
    new user's persistent Zo workspace.
    """
    users = _load_users()
    if users:
        # Check if any admin already exists
        for uid, u in users.items():
            if u.get("tier") == "admin":
                return uid, f"Owner already exists: {u.get('email', uid)}"

    user_id, err = register_user(
        email,
        password,
        name=name,
        tier="admin",
        storage_mode="zo",
    )
    if user_id is None:
        return None, err

    # Migrate existing data into the new workspace
    if migrate_from and migrate_from.exists():
        from vastian.server.worker_pool import ZO_WORKSPACES_DIR

        workspace = ZO_WORKSPACES_DIR / user_id
        workspace.mkdir(parents=True, exist_ok=True)
        for subdir in ["docs", "state", "config", "saved_prompts"]:
            src = migrate_from / subdir
            dst = workspace / subdir
            if src.exists() and src.is_dir():
                if dst.exists():
                    shutil.rmtree(dst)
                shutil.copytree(src, dst)
                logger.info("Migrated %s -> %s", src, dst)

    logger.info("Bootstrapped owner user: %s (%s)", user_id, email)
    return user_id, ""


def user_count() -> int:
    """Return the number of registered users."""
    return len(_load_users())


def make_auth_cookie_header(token: str, request: Request) -> str:
    """Build Set-Cookie header for JWT. Secure only over HTTPS."""
    secure = "; Secure" if request.url.scheme == "https" else ""
    return f"vastian_token={token}; HttpOnly{secure}; SameSite=Lax; Path=/; Max-Age=86400"


def make_clear_cookie_header() -> str:
    """Build Set-Cookie header to clear the auth cookie."""
    return "vastian_token=; Path=/; Max-Age=0"


# ── FastAPI Auth Middleware ───────────────────────────────────────

# Paths that don't require auth
_PUBLIC_PATHS = frozenset(
    {
        "/",
        "/mobile",
        "/login",
        "/api/health",
        "/api/auth/login",
        "/api/auth/register",
        "/api/auth/bootstrap",
        "/api/auth/login-passphrase",
        "/api/auth/bootstrap-passphrase",
        "/api/auth/logout",
        "/api/setup/status",
        "/api/setup",
        "/api/onboard",
        "/api/ollama/status",
        "/api/email/capture",
        "/mcp",
    }
)

# Path prefixes that don't require auth
_PUBLIC_PREFIXES = ("/static/", "/assets/", "/docs/")

# UI paths: unauthenticated GET redirects to /login instead of 401
_UI_PATHS = frozenset(
    {"/dashboard", "/inbox", "/chat", "/metrics", "/scenarios", "/sessions", "/personas", "/boards", "/plans"}
)


def _is_localhost(request: Request) -> bool:
    """Check if the request is from localhost (desktop app / CLI)."""
    client = request.client
    if client is None:
        return False
    host = client.host
    return host in ("127.0.0.1", "::1", "localhost")


class AuthMiddleware(BaseHTTPMiddleware):
    """JWT auth middleware for cloud endpoints.

    - Localhost requests resolve to the **owner** user (first admin).
    - Public paths bypass auth entirely.
    - All other paths require a valid JWT in the Authorization header.
    - On success, sets ``request.state.user_id`` and ``request.state.user``.
    """

    async def dispatch(
        self,
        request: Request,
        call_next: RequestResponseEndpoint,
    ) -> Response:
        path = request.url.path

        # Public paths — no auth needed
        if path in _PUBLIC_PATHS or any(path.startswith(p) for p in _PUBLIC_PREFIXES):
            request.state.user_id = self._resolve_localhost_user()
            request.state.user = None
            return await call_next(request)

        # Localhost bypass — resolve to owner
        if _is_localhost(request):
            owner_id = self._resolve_localhost_user()
            request.state.user_id = owner_id
            request.state.user = get_user(owner_id) if owner_id else None
            return await call_next(request)

        # Extract JWT from Authorization header or Cookie
        token = None
        auth_header = request.headers.get("Authorization", "")
        if auth_header.startswith("Bearer "):
            token = auth_header[7:]
        else:
            # Check Cookie for vastian_token
            cookie_header = request.headers.get("Cookie", "")
            for part in cookie_header.split(";"):
                part = part.strip()
                if part.startswith("vastian_token="):
                    token = part.split("=", 1)[1].strip().strip('"')
                    break

        if not token:
            # GET to UI paths: redirect to login
            if request.method == "GET" and path in _UI_PATHS:
                next_url = request.query_params.get("next", path)
                return RedirectResponse(url=f"/login?next={next_url}", status_code=302)
            return JSONResponse(
                {"ok": False, "error": "Authentication required"},
                status_code=401,
            )

        # token is set from header or cookie
        payload = verify_jwt(token)
        if payload is None:
            if request.method == "GET" and path in _UI_PATHS:
                next_url = request.query_params.get("next", path)
                return RedirectResponse(url=f"/login?next={next_url}", status_code=302)
            return JSONResponse(
                {"ok": False, "error": "Invalid or expired token"},
                status_code=401,
            )

        user_id = payload.get("sub", "")
        user = get_user(user_id)
        if user is None:
            if request.method == "GET" and path in _UI_PATHS:
                next_url = request.query_params.get("next", path)
                return RedirectResponse(url=f"/login?next={next_url}", status_code=302)
            return JSONResponse(
                {"ok": False, "error": "User not found"},
                status_code=401,
            )

        request.state.user_id = user_id
        request.state.user = user
        return await call_next(request)

    @staticmethod
    def _resolve_localhost_user() -> str:
        """Resolve the user ID for localhost requests.

        Returns the owner's user_id if one exists, otherwise "unregistered"
        (the setup wizard or bootstrap flow will handle registration).
        """
        owner_id, _ = get_owner_user()
        return owner_id or "unregistered"
