from __future__ import annotations

import hashlib
import hmac
import json
import os
import secrets
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from threading import RLock
from typing import Dict, Iterable, List, Optional, Set


class AuthorizationError(PermissionError):
    """Raised when authentication or authorization fails."""


ALL_PERMISSIONS = (
    "model:view",
    "model:import",
    "model:promote",
    "model:rollback",
    "audit:view",
    "admin:user_manage",
    "admin:role_manage",
)


DEFAULT_ROLES = {
    "admin": {
        "description": "Full access to model lifecycle and RBAC operations.",
        "permissions": list(ALL_PERMISSIONS),
    },
    "operator": {
        "description": "Operational role for import/promote/rollback and audits.",
        "permissions": [
            "model:view",
            "model:import",
            "model:promote",
            "model:rollback",
            "audit:view",
        ],
    },
    "viewer": {
        "description": "Read-only role for model and audit visibility.",
        "permissions": ["model:view", "audit:view"],
    },
}


def _utcnow_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _normalize_permissions(permissions: Iterable[str]) -> List[str]:
    return sorted({str(item).strip() for item in permissions if str(item).strip()})


@dataclass(frozen=True)
class AuthenticatedUser:
    username: str
    roles: List[str]
    permissions: List[str]


class RBACStore:
    """Simple JSON-backed RBAC store with salted-password authentication."""

    def __init__(
        self,
        store_path: str | Path,
        *,
        bootstrap_admin_username: str = "admin",
        bootstrap_admin_password: Optional[str] = None,
    ) -> None:
        self.store_path = Path(store_path)
        self.store_path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = RLock()
        self.bootstrap_admin_username = bootstrap_admin_username.strip() or "admin"
        self.bootstrap_admin_password = (
            bootstrap_admin_password
            or os.environ.get("MODEL_MANAGER_ADMIN_PASSWORD")
            or "admin123!"
        )
        self._ensure_bootstrap_state()

    def _load(self) -> Dict[str, Dict[str, dict]]:
        if not self.store_path.exists():
            return {"roles": {}, "users": {}}
        with self.store_path.open("r", encoding="utf-8") as fh:
            payload = json.load(fh)
        if not isinstance(payload, dict):
            return {"roles": {}, "users": {}}
        roles = payload.get("roles")
        users = payload.get("users")
        return {
            "roles": dict(roles) if isinstance(roles, dict) else {},
            "users": dict(users) if isinstance(users, dict) else {},
        }

    def _save(self, payload: Dict[str, Dict[str, dict]]) -> None:
        with self.store_path.open("w", encoding="utf-8") as fh:
            json.dump(payload, fh, indent=2, ensure_ascii=True)

    @staticmethod
    def _hash_password(password: str, salt: str) -> str:
        raw = f"{salt}:{password}".encode("utf-8")
        return hashlib.sha256(raw).hexdigest()

    @staticmethod
    def _make_user_record(password: str, roles: Iterable[str]) -> Dict[str, object]:
        salt = secrets.token_hex(16)
        return {
            "salt": salt,
            "password_hash": RBACStore._hash_password(password, salt),
            "roles": sorted({str(role).strip() for role in roles if str(role).strip()}),
            "active": True,
            "created_at": _utcnow_iso(),
            "updated_at": _utcnow_iso(),
        }

    def _ensure_bootstrap_state(self) -> None:
        with self._lock:
            payload = self._load()
            changed = False

            roles = payload.setdefault("roles", {})
            users = payload.setdefault("users", {})

            for role_name, role_cfg in DEFAULT_ROLES.items():
                existing = roles.get(role_name)
                if not isinstance(existing, dict):
                    roles[role_name] = {
                        "description": str(role_cfg.get("description", "")),
                        "permissions": _normalize_permissions(role_cfg.get("permissions", [])),
                    }
                    changed = True
                    continue
                normalized = {
                    "description": str(
                        existing.get("description", role_cfg.get("description", ""))
                    ),
                    "permissions": _normalize_permissions(
                        existing.get("permissions", role_cfg.get("permissions", []))
                    ),
                }
                if normalized != existing:
                    roles[role_name] = normalized
                    changed = True

            if self.bootstrap_admin_username not in users:
                users[self.bootstrap_admin_username] = self._make_user_record(
                    self.bootstrap_admin_password,
                    roles=["admin"],
                )
                changed = True
            else:
                admin_record = users[self.bootstrap_admin_username]
                if not isinstance(admin_record, dict):
                    users[self.bootstrap_admin_username] = self._make_user_record(
                        self.bootstrap_admin_password,
                        roles=["admin"],
                    )
                    changed = True
                else:
                    role_set = {
                        str(item).strip()
                        for item in admin_record.get("roles", [])
                        if str(item).strip()
                    }
                    if "admin" not in role_set:
                        role_set.add("admin")
                        admin_record["roles"] = sorted(role_set)
                        admin_record["updated_at"] = _utcnow_iso()
                        changed = True

            if changed:
                self._save(payload)

    def _build_user(self, username: str, payload: Dict[str, Dict[str, dict]]) -> AuthenticatedUser:
        users = payload.get("users", {})
        user_record = users.get(username)
        if not isinstance(user_record, dict):
            raise AuthorizationError("User not found.")
        if not bool(user_record.get("active", True)):
            raise AuthorizationError("User is disabled.")

        role_names = [
            str(item).strip()
            for item in user_record.get("roles", [])
            if str(item).strip()
        ]
        permissions = self._permissions_for_roles(role_names, payload)
        return AuthenticatedUser(
            username=username,
            roles=sorted(set(role_names)),
            permissions=sorted(permissions),
        )

    @staticmethod
    def _permissions_for_roles(role_names: Iterable[str], payload: Dict[str, Dict[str, dict]]) -> Set[str]:
        roles = payload.get("roles", {})
        permissions: Set[str] = set()
        for role_name in role_names:
            role_cfg = roles.get(role_name)
            if not isinstance(role_cfg, dict):
                continue
            permissions.update(_normalize_permissions(role_cfg.get("permissions", [])))
        return permissions

    def authenticate(self, username: str, password: str) -> AuthenticatedUser:
        username = (username or "").strip()
        if not username or not password:
            raise AuthorizationError("Username and password are required.")

        with self._lock:
            payload = self._load()
            user_record = payload.get("users", {}).get(username)
            if not isinstance(user_record, dict):
                raise AuthorizationError("Invalid username or password.")
            if not bool(user_record.get("active", True)):
                raise AuthorizationError("User is disabled.")

            salt = str(user_record.get("salt", ""))
            expected_hash = str(user_record.get("password_hash", ""))
            actual_hash = self._hash_password(password, salt)
            if not hmac.compare_digest(actual_hash, expected_hash):
                raise AuthorizationError("Invalid username or password.")
            return self._build_user(username, payload)

    def get_user(self, username: str) -> AuthenticatedUser:
        with self._lock:
            payload = self._load()
            return self._build_user(username, payload)

    def has_permission(self, username: str, permission: str) -> bool:
        user = self.get_user(username)
        return permission in user.permissions

    def require_permission(self, username: str, permission: str) -> None:
        if not self.has_permission(username, permission):
            raise AuthorizationError(f"Permission denied: {permission}")

    def list_users(self) -> List[Dict[str, object]]:
        with self._lock:
            payload = self._load()
            users = payload.get("users", {})
            result: List[Dict[str, object]] = []
            for username in sorted(users.keys()):
                record = users.get(username)
                if not isinstance(record, dict):
                    continue
                result.append(
                    {
                        "username": username,
                        "roles": sorted(
                            {
                                str(role).strip()
                                for role in record.get("roles", [])
                                if str(role).strip()
                            }
                        ),
                        "active": bool(record.get("active", True)),
                        "created_at": record.get("created_at"),
                        "updated_at": record.get("updated_at"),
                    }
                )
            return result

    def list_roles(self) -> Dict[str, Dict[str, object]]:
        with self._lock:
            payload = self._load()
            roles = payload.get("roles", {})
            normalized: Dict[str, Dict[str, object]] = {}
            for role_name, role_cfg in sorted(roles.items()):
                if not isinstance(role_cfg, dict):
                    continue
                normalized[role_name] = {
                    "description": str(role_cfg.get("description", "")),
                    "permissions": _normalize_permissions(role_cfg.get("permissions", [])),
                }
            return normalized

    def create_user(self, username: str, password: str, roles: Iterable[str]) -> Dict[str, object]:
        username = (username or "").strip()
        if not username:
            raise ValueError("username is required")
        if not password:
            raise ValueError("password is required")

        role_set = sorted({str(role).strip() for role in roles if str(role).strip()})
        if not role_set:
            raise ValueError("at least one role is required")

        with self._lock:
            payload = self._load()
            existing_users = payload.setdefault("users", {})
            existing_roles = payload.setdefault("roles", {})
            missing = [role for role in role_set if role not in existing_roles]
            if missing:
                raise ValueError(f"unknown roles: {missing}")
            if username in existing_users:
                raise ValueError(f"user already exists: {username}")
            existing_users[username] = self._make_user_record(password, role_set)
            self._save(payload)
            return {"username": username, "roles": role_set, "active": True}

    def set_user_roles(self, username: str, roles: Iterable[str]) -> Dict[str, object]:
        username = (username or "").strip()
        if not username:
            raise ValueError("username is required")
        role_set = sorted({str(role).strip() for role in roles if str(role).strip()})
        if not role_set:
            raise ValueError("at least one role is required")

        with self._lock:
            payload = self._load()
            users = payload.setdefault("users", {})
            roles_cfg = payload.setdefault("roles", {})
            if username not in users or not isinstance(users[username], dict):
                raise ValueError(f"user not found: {username}")
            missing = [role for role in role_set if role not in roles_cfg]
            if missing:
                raise ValueError(f"unknown roles: {missing}")
            users[username]["roles"] = role_set
            users[username]["updated_at"] = _utcnow_iso()
            self._save(payload)
            return {"username": username, "roles": role_set}

    def create_role(
        self,
        role_name: str,
        permissions: Iterable[str],
        *,
        description: str = "",
    ) -> Dict[str, object]:
        role_name = (role_name or "").strip()
        if not role_name:
            raise ValueError("role_name is required")
        permission_list = _normalize_permissions(permissions)
        if not permission_list:
            raise ValueError("permissions cannot be empty")

        with self._lock:
            payload = self._load()
            roles = payload.setdefault("roles", {})
            if role_name in roles:
                raise ValueError(f"role already exists: {role_name}")
            roles[role_name] = {
                "description": str(description or ""),
                "permissions": permission_list,
            }
            self._save(payload)
            return {
                "role_name": role_name,
                "description": str(description or ""),
                "permissions": permission_list,
            }

