import asyncio
import logging
import uuid
import jwt
import base64
import hmac
import hashlib
import requests
import os
from pathlib import Path


from datetime import datetime, timedelta
import pytz
from pytz import UTC
from typing import Optional, Union, List, Dict

from open_webui.models.users import User, Users

from open_webui.constants import ERROR_MESSAGES
from open_webui.env import (
    ENABLE_OTEL,
    WEBUI_SECRET_KEY,
    TRUSTED_SIGNATURE_KEY,
    STATIC_DIR,
    SRC_LOG_LEVELS,
    GLOBAL_LOG_LEVEL,
    REDIS_KEY_PREFIX,
    WEBUI_AUTH_TRUSTED_EMAIL_HEADER,
    is_superman_admin,
)

from fastapi import BackgroundTasks, Depends, HTTPException, Request, Response, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from passlib.context import CryptContext


logging.getLogger("passlib").setLevel(logging.ERROR)

log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["OAUTH"])

SESSION_SECRET = WEBUI_SECRET_KEY
ALGORITHM = "HS256"

# Role whitelist used by `get_verified_user`, `get_admin_user`, and
# `get_verified_user_by_token`. Centralizing prevents drift between
# HTTP and WebSocket auth paths (Bug #F).
VERIFIED_USER_ROLES = {"user", "admin"}

##############
# Auth Utils
##############


def verify_signature(payload: str, signature: str) -> bool:
    """
    Verifies the HMAC signature of the received payload.
    """
    try:
        expected_signature = base64.b64encode(
            hmac.new(TRUSTED_SIGNATURE_KEY, payload.encode(), hashlib.sha256).digest()
        ).decode()

        # Compare securely to prevent timing attacks
        return hmac.compare_digest(expected_signature, signature)

    except Exception:
        return False


def override_static(path: str, content: str):
    """
    安全地覆盖静态文件

    Args:
        path: 相对于 STATIC_DIR 的文件路径
        content: Base64 编码的文件内容

    Returns:
        bool: 操作是否成功
    """
    try:
        # 将路径转换为 Path 对象并规范化
        static_dir = Path(STATIC_DIR).resolve()
        target_path = (static_dir / path).resolve()

        # 安全检查：确保目标路径在 STATIC_DIR 内
        if not target_path.is_relative_to(static_dir):
            log.error(f"Path traversal attempt detected: {path}")
            log.error(f"Resolved path {target_path} is outside {static_dir}")
            return False

        # 创建父目录
        target_path.parent.mkdir(parents=True, exist_ok=True)

        # 写入文件（二进制模式）
        target_path.write_bytes(base64.b64decode(content))

        log.info(f"Successfully wrote static file: {target_path.relative_to(static_dir)}")
        return True

    except (ValueError, OSError) as e:
        log.error(f"Failed to write static file {path}: {str(e)}")
        return False
    except base64.binascii.Error as e:
        log.error(f"Invalid base64 content for {path}: {str(e)}")
        return False


def get_license_data(app, key):
    """
    从许可证服务器获取许可证数据

    Args:
        app: FastAPI 应用实例
        key: 许可证密钥

    Returns:
        bool: 是否成功获取许可证数据
    """
    # 检查是否跳过许可证验证（开发/测试环境）
    # 注意：不建议在生产环境使用此选项
    skip_license_check = os.environ.get("SKIP_LICENSE_CHECK", "").lower() == "true"

    if skip_license_check:
        log.warning("SKIP_LICENSE_CHECK is enabled: License validation bypassed")
        log.warning("This should ONLY be used in development/testing environments")
        return True

    if key:
        try:
            # 通过许可证服务器获取授权的配置信息
            res = requests.post(
                "https://ai.lizhisec.com/api/v1/license/",
                json={"key": key, "version": "1"},
                timeout=5,
                verify=True,  # 显式启用 TLS 证书验证
            )

            if getattr(res, "ok", False):
                payload = getattr(res, "json", lambda: {})()
                for k, v in payload.items():
                    if k == "resources":
                        for p, c in v.items():
                            override_static(p, c)
                    elif k == "count":
                        setattr(app.state, "USER_COUNT", v)
                    elif k == "name":
                        setattr(app.state, "WEBUI_NAME", v)
                    elif k == "metadata":
                        setattr(app.state, "LICENSE_METADATA", v)
                return True
            else:
                log.error(
                    f"License: retrieval issue: {getattr(res, 'text', 'unknown error')}"
                )
        except requests.exceptions.SSLError as ex:
            log.exception(f"License: SSL Certificate verification failed: {ex}")
        except requests.exceptions.Timeout as ex:
            log.exception(f"License: Request timeout: {ex}")
        except requests.exceptions.RequestException as ex:
            log.exception(f"License: Request failed: {ex}")
        except Exception as ex:
            log.exception(f"License: Uncaught Exception: {ex}")
    return False


bearer_security = HTTPBearer(auto_error=False)
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")


async def verify_password(plain_password, hashed_password):
    """异步校验密码，避免 bcrypt CPU 哈希阻塞事件循环。

    始终返回 bool（与 open-webui-0.11.0 reference 一致）。调用方若写
    `if not await verify_password(...)`，None 会被判为 False，
    需始终保持返回 bool 类型。
    """
    if not hashed_password:
        return False
    return await asyncio.to_thread(pwd_context.verify, plain_password, hashed_password)


async def get_password_hash(password):
    """异步生成密码哈希，避免 bcrypt CPU 哈希阻塞事件循环。"""
    return await asyncio.to_thread(pwd_context.hash, password)


def create_token(data: dict, expires_delta: Union[timedelta, None] = None) -> str:
    payload = data.copy()

    if expires_delta:
        expire = datetime.now(UTC) + expires_delta
        payload.update({"exp": expire})

    # jti: 用于按 token 撤销（登出）；iat: 用于按用户 revoked_at 撤销（OIDC back-channel logout）
    payload.update({"jti": str(uuid.uuid4()), "iat": datetime.now(UTC)})

    encoded_jwt = jwt.encode(payload, SESSION_SECRET, algorithm=ALGORITHM)
    return encoded_jwt


def decode_token(token: str) -> Optional[dict]:
    """
    解码 JWT token

    Args:
        token: JWT token 字符串

    Returns:
        解码后的字典，失败时返回 None
    """
    try:
        decoded = jwt.decode(token, SESSION_SECRET, algorithms=[ALGORITHM])
        return decoded
    except (jwt.ExpiredSignatureError, jwt.InvalidTokenError, jwt.DecodeError) as e:
        log.debug(f"Token decode failed: {type(e).__name__}")
        return None


async def is_valid_token(decoded: dict, redis=None) -> bool:
    """检查 JWT 是否已被撤销。两种机制：
    1. 按 token (jti) —— 用户主动登出（已知 jti）。
    2. 按用户 (revoked_at) —— OIDC back-channel logout，jti 未知时拒绝
       iat <= revoked_at 的所有 token。
    无 redis 时视为有效（向后兼容）。
    """
    if redis:
        jti = decoded.get("jti")
        if jti:
            revoked = await redis.get(f"{REDIS_KEY_PREFIX}:auth:token:{jti}:revoked")
            if revoked:
                return False

        user_id = decoded.get("id")
        if user_id:
            revoked_at = await redis.get(
                f"{REDIS_KEY_PREFIX}:auth:user:{user_id}:revoked_at"
            )
            if revoked_at:
                try:
                    revoked_at_ts = int(revoked_at)
                    token_iat = decoded.get("iat")
                    # 无 iat 视为旧 token，无法校验签发时间，拒绝
                    if token_iat is None or token_iat <= revoked_at_ts:
                        return False
                except (ValueError, TypeError):
                    pass
    return True


async def invalidate_token(request, token: str) -> None:
    """将 token 标记为已撤销（登出时调用），写入 redis 并设置 TTL。"""
    decoded = decode_token(token)
    if not decoded:
        return

    redis = getattr(request.app.state, "redis", None)
    if redis:
        jti = decoded.get("jti")
        exp = decoded.get("exp")
        if jti and exp:
            ttl = exp - int(datetime.now(UTC).timestamp())
            if ttl > 0:
                await redis.set(
                    f"{REDIS_KEY_PREFIX}:auth:token:{jti}:revoked",
                    "1",
                    ex=ttl,
                )


def extract_token_from_auth_header(auth_header: Optional[str]) -> Optional[str]:
    """
    从 Authorization header 中提取 token（Bearer scheme）。

    Args:
        auth_header: Authorization header 字符串

    Returns:
        提取的 token；header 缺失或前缀不匹配时返回 None。
    """
    if not auth_header or not auth_header.startswith("Bearer "):
        return None
    return auth_header[len("Bearer ") :]


def create_api_key() -> str:
    """
    生成新的 API key

    Returns:
        格式为 sk-{uuid} 的 API key
    """
    key = str(uuid.uuid4()).replace("-", "")
    return f"sk-{key}"


def get_http_authorization_cred(auth_header: Optional[str]) -> Optional[HTTPAuthorizationCredentials]:
    """
    从 Authorization header 获取认证凭据

    Args:
        auth_header: Authorization header 字符串

    Returns:
        HTTPAuthorizationCredentials 对象或 None
    """
    if not auth_header:
        return None
    try:
        scheme, credentials = auth_header.split(" ", 1)
        cred = HTTPAuthorizationCredentials(scheme=scheme, credentials=credentials)
        return cred
    except (ValueError, AttributeError) as e:
        log.debug(f"Failed to parse auth header: {e}")
        return None


async def get_current_user(
    request: Request,
    response: Response,
    background_tasks: BackgroundTasks,
    auth_token: HTTPAuthorizationCredentials = Depends(bearer_security),
):
    token = None

    if auth_token is not None:
        token = auth_token.credentials

    if token is None and "token" in request.cookies:
        token = request.cookies.get("token")

    # Fallback to request.state.token (set by AuthTokenMiddleware,
    # e.g. for x-api-key header). Without this, Depends(get_current_user)
    # cannot serve the x-api-key credential path.
    if token is None and getattr(request.state, "token", None):
        token = request.state.token.credentials

    if token is None:
        raise HTTPException(status_code=403, detail="Not authenticated")

    # auth by api key
    if token.startswith("sk-"):
        # `enable_api_key` is normally written by AuthTokenMiddleware (see
        # utils/asgi_middleware.py:113-115). Use a safe default of True to
        # avoid AttributeError when the middleware is not on the path
        # (e.g., tests, custom sub-apps, or external mounts).
        if not getattr(request.state, "enable_api_key", True):
            raise HTTPException(
                status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.API_KEY_NOT_ALLOWED
            )

        # Resolve user via API key (also re-checks enable_api_key, permissions,
        # and endpoint restrictions inside get_current_user_by_api_key — see
        # open-webui-0.11.0 auth.py:433-488).
        user = await get_current_user_by_api_key(request, token)

        # Bug #C: expose the resolved user to outer middleware (audit,
        # rate-limit, etc.) via request scope — mirrors v0.11.0 behavior.
        request.state.user = user

        # Bug #E: OTel span attributes
        if ENABLE_OTEL:
            from opentelemetry import trace

            current_span = trace.get_current_span()
            if current_span:
                current_span.set_attribute("client.user.id", user.id)
                current_span.set_attribute("client.user.email", user.email)
                current_span.set_attribute("client.user.role", user.role)
                current_span.set_attribute("client.auth.type", "api_key")
        return user

    # auth by jwt token
    try:
        try:
            data = decode_token(token)
        except Exception as e:
            log.debug("decode_token raised unexpectedly: %s", e)
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Invalid token",
            )

        if data is not None and "id" in data:
            if not await is_valid_token(data, getattr(request.app.state, "redis", None)):
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail=ERROR_MESSAGES.INVALID_TOKEN,
                )

            user = await Users.get_user_by_id(data["id"])
            if user is None:
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail=ERROR_MESSAGES.INVALID_TOKEN,
                )
            else:
                # Bug #B: when a reverse proxy asserts the authenticated email
                # via a trusted header, ensure the JWT subject matches.
                # Without this check an attacker could present any valid JWT
                # while the proxy asserts an unrelated identity.
                if WEBUI_AUTH_TRUSTED_EMAIL_HEADER:
                    trusted_email = request.headers.get(
                        WEBUI_AUTH_TRUSTED_EMAIL_HEADER, ""
                    ).lower()
                    if trusted_email and user.email != trusted_email:
                        raise HTTPException(
                            status_code=status.HTTP_401_UNAUTHORIZED,
                            detail="User mismatch. Please sign in again.",
                        )

                # Bug #C: expose resolved user to outer middleware.
                request.state.user = user

                # Bug #E: OTel span attributes
                if ENABLE_OTEL:
                    from opentelemetry import trace

                    current_span = trace.get_current_span()
                    if current_span:
                        current_span.set_attribute("client.user.id", user.id)
                        current_span.set_attribute("client.user.email", user.email)
                        current_span.set_attribute("client.user.role", user.role)
                        current_span.set_attribute("client.auth.type", "jwt")

                # Refresh the user's last active timestamp via fire-and-forget
                # asyncio task (BackgroundTasks.add_task does not await
                # coroutines — see open-webui-0.11.0 auth.py:407-408).
                asyncio.create_task(Users.update_user_last_active_by_id(user.id))

            return user
        else:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.UNAUTHORIZED,
            )
    except Exception as e:
        # Bug #A: scrub stale auth cookies so the FE does not loop on 401
        # after logout / revocation / disabled account. Mirror v0.11.0
        # auth.py:418-430. Cookie deletes are best-effort and do not mask
        # the original exception.
        if request.cookies.get("token"):
            response.delete_cookie("token")
        if request.cookies.get("oauth_id_token"):
            response.delete_cookie("oauth_id_token")
        if request.cookies.get("oauth_session_id"):
            response.delete_cookie("oauth_session_id")
        raise e


def _is_taiji_app_enabled_for_auth() -> bool:
    """Taiji 应用已加载且启用时才尝试员工 coding_gateway_api_key 鉴权。"""
    try:
        from open_webui.apps import app_loader

        taiji = app_loader.apps.get("taiji")
        return taiji is not None and taiji.enabled
    except Exception:
        return False


async def _try_resolve_user_from_coding_gateway_api_key(api_key: str) -> Optional[User]:
    """
    在 Taiji 可用时用员工 coding_gateway_api_key 解析用户；
    应用未加载/未启用或依赖异常时返回 None，不影响平台 api_key 主流程。
    """
    if not _is_taiji_app_enabled_for_auth():
        return None
    try:
        from open_webui.apps.taiji.utils.coding_api_key_auth import (
            resolve_user_from_coding_gateway_api_key,
        )

        return await resolve_user_from_coding_gateway_api_key(api_key)
    except (ImportError, ModuleNotFoundError):
        return None
    except Exception as e:
        log.debug("coding gateway api key resolve skipped: %s", e)
        return None


async def get_current_user_by_api_key(
    request: Request,
    api_key: str,
):
    # Each function call manages its own short-lived session internally.
    user = await Users.get_user_by_api_key(api_key)
    if user is None:
        user = await _try_resolve_user_from_coding_gateway_api_key(api_key)

    if user is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.INVALID_TOKEN,
        )

    # Global ENABLE_API_KEY + optional endpoint allow-list.
    # 不检查 permissions.features.api_keys：只要全局开关开启且 key 有效即可。
    config = getattr(request.app.state, "config", None)
    if config is None:
        enable_api_keys = True
    else:
        enable_api_keys = getattr(config, "ENABLE_API_KEY", True)
    if not enable_api_keys:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail=ERROR_MESSAGES.API_KEY_NOT_ALLOWED,
        )

    if config is not None and getattr(config, "ENABLE_API_KEY_ENDPOINT_RESTRICTIONS", False):
        allowed_endpoints = str(
            getattr(config, "API_KEY_ALLOWED_ENDPOINTS", "")
        )
        allowed_paths = [
            p.strip() for p in allowed_endpoints.split(",") if p.strip()
        ]
        request_path = (
            request.scope.get("path")
            if isinstance(request.scope, dict) and request.scope.get("path") is not None
            else request.url.path
        )
        is_allowed = any(
            request_path == allowed or request_path.startswith(allowed + "/")
            for allowed in allowed_paths
        )
        if not is_allowed:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
            )

    # Fire-and-forget last_active update.
    asyncio.create_task(Users.update_user_last_active_by_id(user.id))
    return user


def get_verified_user(user=Depends(get_current_user)):
    # Use getattr for mock safety — matches user_is_admin / get_verified_user_by_token
    # style. UserModel.role defaults to "pending" (models/users.py:49), so the
    # production path is unaffected; non-UserModel mocks (e.g. plain dicts)
    # won't crash with AttributeError here.
    if getattr(user, "role", "pending") not in VERIFIED_USER_ROLES:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
    return user


def user_is_admin(user) -> bool:
    """
    判断当前用户是否为管理员（与 get_verified_user / get_admin_user 的角色约定一致）。
    用于同一路由内需区分管理员与普通用户权限的场景；仅允许管理员访问请用 Depends(get_admin_user)。
    """
    return getattr(user, "role", "user") == "admin"


def get_admin_user(user=Depends(get_current_user)):
    if user.role != "admin":
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
    return user


def get_super_user(user=Depends(get_current_user)):
    if is_superman_admin(user):
        return user
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )


def is_super_user(user=Depends(get_current_user)):
    return is_superman_admin(user)


async def get_verified_user_by_token(token: str, redis=None):
    """从 WebSocket token 解析已验证用户（async DB，不阻塞事件循环）。

    用于 Socket.IO connect/user-join 等脱离 FastAPI HTTP Depends 链的场景。
    redis 参数支持 token 撤销校验（与 get_current_user 一致）。
    """
    decoded = decode_token(token)
    if decoded is None or "id" not in decoded:
        return None

    if not await is_valid_token(decoded, redis):
        return None

    user = await Users.get_user_by_id(decoded["id"])
    if user is None:
        return None

    if getattr(user, "role", "user") not in VERIFIED_USER_ROLES:
        return None

    return user