"""
FastAPI版本的mwauth认证模块
提供基于api gateway kong的JWT认证和cookie认证
"""

import json
import hashlib
from typing import Optional, Dict, Any, Callable, Union
from datetime import timedelta, datetime
from uuid import uuid4
from redis import Redis
from fastapi import Request, Response, HTTPException, Depends, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.responses import JSONResponse
import inspect
from .base_auth import User

class FastApiRedisSession:
    """
    FastAPI Redis 会话管理器
    """
    def __init__(self, redis_client: Redis, prefix: str = 'session'):
        assert redis_client is not None, 'redis 必须要指定'
        self.redis = redis_client
        self.prefix = prefix

    def _get_sid(self, sid: str) -> str:
        if self.prefix:
            return f"{self.prefix}:{sid}"
        return sid

    def generate_sid(self) -> str:
        return str(uuid4())

    def get_session(self, sid: Optional[str]) -> Dict[str, Any]:
        if not sid:
            sid = self.generate_sid()
            return {'sid': sid, 'new': True, '_data': {}}
        
        val = self.redis.get(self._get_sid(sid))
        if val is not None:
            try:
                data = json.loads(val.decode())
                return {'sid': sid, 'new': False, '_data': data}
            except json.JSONDecodeError:
                return {'sid': sid, 'new': True, '_data': {}}
        return {'sid': sid, 'new': True, '_data': {}}

    def save_session(self, sid: str, data: Dict[str, Any], expiration_time: Optional[int] = None):
        """保存会话到Redis"""
        if expiration_time is None:
            # 默认2天过期
            expiration_time = 2 * 24 * 3600  # 2 days in seconds
        
        session_data = {k: v for k, v in data.items() if k != 'sid' and not k.startswith('_')}
        val = json.dumps(session_data)
        
        self.redis.setex(
            self._get_sid(sid),
            expiration_time,
            val
        )

    def delete_session(self, sid: str):
        """删除会话"""
        self.redis.delete(self._get_sid(sid))

    def get_expiration_time(self, permanent: bool = False, temp_expiration_time: int = 0) -> int:
        """获取会话过期时间（秒）"""
        if permanent:
            return 30 * 24 * 3600  # 30天
        if temp_expiration_time > 0:
            return temp_expiration_time
        return 2 * 24 * 3600  # 2天


class FastApiSessionMiddleware(BaseHTTPMiddleware):
    """
    FastAPI会话中间件
    """
    def __init__(self, dispatch: Optional[Callable] = None, redis_client: Redis = None, prefix: str = 'session'):
        super().__init__(dispatch)
        self.session_manager = FastApiRedisSession(redis_client, prefix)

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
        # 提取session ID
        sessionid = request.cookies.get('sessionid')
        
        # 检查是否通过Authorization头部传递了JWT（用于Kong认证）
        auth_header = request.headers.get('authorization', None)
        jwt_token = None
        if auth_header and auth_header.lower().startswith('bearer '):
            jwt_token = auth_header[7:]
        
        # 如果没有cookie中的sessionid，但有JWT，使用JWT作为session ID
        session_id_to_use = sessionid or jwt_token
        
        # 获取会话数据
        session_data = self.session_manager.get_session(session_id_to_use)
        request.state.session = session_data['_data'].copy() if '_data' in session_data else {}
        request.state.session_sid = session_data.get('sid', session_id_to_use)
        request.state.session_new = session_data.get('new', True)
        request.state.session_modified = False
        
        # 处理请求
        response = await call_next(request)
        
        # 如果会话被修改，保存到Redis
        if hasattr(request.state, 'session_modified') and request.state.session_modified:
            if hasattr(request.state, 'session'):
                # 更新temp_expiration_time字段如果存在
                temp_exp_time = request.state.session.get('temp_expiration_time', 0)
                
                # 从会话中移除内部字段
                session_copy = request.state.session.copy()
                session_copy.pop('temp_expiration_time', None)  # 移除临时过期时间字段
                
                self.session_manager.save_session(
                    request.state.session_sid,
                    session_copy,
                    self.session_manager.get_expiration_time(temp_expiration_time=temp_exp_time)
                )
                
                # 设置cookie
                if not response.headers.get('set-cookie') and sessionid:
                    # 如果响应中没有设置cookie，我们设置它（只有在原始请求有cookie时）
                    response.set_cookie(
                        'sessionid',
                        request.state.session_sid,
                        max_age=2 * 24 * 3600,  # 2天
                        httponly=True
                    )
        
        return response


class CurrentUser:
    """
    当前用户信息存储类
    """
    def __init__(self):
        self.user_name: Optional[str] = None
        self.user_id: Optional[str] = None
        self.jwt: Optional[str] = None
        self.current_user: Optional[User] = None


class FastApiKongAuth:
    """
    FastAPI 版本的 KongAuth 认证类
    提供基于api gateway kong的JWT认证和cookie认证
    认证失败时抛出401异常
    """
    def __init__(self, redis_client: Redis, session_prefix: str = 'session'):
        self.session_manager = FastApiRedisSession(redis_client, session_prefix)
        self.current_user = CurrentUser()

    def _set_auth(self, request: Request):
        """
        保存用户，jwt等auth info
        """
        # redis 中session的过期时间，为0时，为第一次设定的时间
        session_expiration_time = getattr(request.app, 'config', {}).get('SESSION_EXPIRATION_TIME', 0)
        
        # 检查是否为开发模式
        development_mode = getattr(request.app, 'config', {}).get('DEVELOPMENT', False)
        if development_mode:
            login_user_name = getattr(request.app, 'config', {}).get('LOGIN_USER_NAME')
            login_user_id = getattr(request.app, 'config', {}).get('LOGIN_USER_ID')
            
            if login_user_name and login_user_id:
                self.current_user.user_name = login_user_name
                self.current_user.user_id = login_user_id
                self.current_user.current_user = User(
                    uid=getattr(request.app, 'config', {}).get('LOGIN_USER_ID'),
                    uname=getattr(request.app, 'config', {}).get('LOGIN_USER_NAME'),
                    systemuser=getattr(request.app, 'config', {}).get('LOGIN_USER_SYSTEMUSER', False),
                    manageuser=getattr(request.app, 'config', {}).get('LOGIN_USER_MANAGEUSER', False),
                    manageuserid=getattr(request.app, 'config', {}).get('LOGIN_USER_MANAGEUSER_ID'),
                    companyid=getattr(request.app, 'config', {}).get('LOGIN_USER_COMPANYID', ''),
                    phone=getattr(request.app, 'config', {}).get('LOGIN_USER_PHONE'),
                    type=request.state.session.get('type', 'appuser') if hasattr(request.state, 'session') else 'appuser'
                )
                self.current_user.jwt = request.state.session_sid if hasattr(request.state, 'session_sid') else None
                
                # 只有session认证才需要设定这个redis的有效期
                if hasattr(request.state, 'session'):
                    request.state.session['temp_expiration_time'] = session_expiration_time
                    request.state.session.update({
                        'uid': self.current_user.current_user.uid,
                        'uname': self.current_user.current_user.uname,
                        'systemuser': self.current_user.current_user.systemuser,
                        'manageuser': self.current_user.current_user.manageuser,
                        'manageuserid': self.current_user.current_user.manageuserid
                    })
                    request.state.session_modified = True
                
                return

        # 通过了jwt 或 key 认证的 api 一定会回传username和userid，否则视为kong的认证不成功
        self.current_user.user_name = request.headers.get('X-Consumer-Username')
        self.current_user.user_id = request.headers.get('X-Consumer-Custom-Id')
        
        # key auth 或jwt auth
        jwt_param = (request.query_params.get('jwt') or 
                    request.query_params.get('apikey') or 
                    request.query_params.get('sessionid') or 
                    request.query_params.get('token'))
        
        self.current_user.jwt = jwt_param
        
        if self.current_user.jwt is None:
            # jwt header authorization: bearer jwt...
            auth_header = request.headers.get('authorization', None)
            if auth_header:
                # 去掉前面的bearer
                if auth_header.lower().startswith('bearer '):
                    self.current_user.jwt = auth_header[7:]
                else:
                    self.current_user.jwt = auth_header
            elif hasattr(request.state, 'session') and request.state.session:
                self.current_user.jwt = request.state.session_sid
                # session 认证可能不经过kong，但也需保证auth 成功，确保代码兼容
                self.current_user.user_name = request.state.session.get('uname', self.current_user.user_name)
                self.current_user.user_id = request.state.session.get('uid', self.current_user.user_id)
            else:
                # header，query，和cookie中都没有 jwt or sessionid or token or apikey，需要重新认证
                self.current_user.user_name = None
                self.current_user.user_id = None
                return
            
            # 只有session 认证才需要设定这个redis的有效期
            if hasattr(request.state, 'session'):
                request.state.session['temp_expiration_time'] = session_expiration_time
                request.state.session_modified = True
        elif (hasattr(request.state, 'session') and 
              request.state.session and 
              request.state.session_sid != self.current_user.jwt):
            # 如果有传jwt ，就使用jwt的，cookie中的session可能是上一次的，需刷新为jwt的
            request.state.session_sid = self.current_user.jwt
            mysession = self.session_manager.get_session(self.current_user.jwt)
            # 只有session 认证才需要设定这个redis的有效期
            if hasattr(request.state, 'session'):
                request.state.session['temp_expiration_time'] = session_expiration_time
                request.state.session.update(mysession.get('_data', {}))
                request.state.session_modified = True
                # 更新用户信息
                self.current_user.user_name = mysession.get('_data', {}).get('uname', self.current_user.user_name)
                self.current_user.user_id = mysession.get('_data', {}).get('uid', self.current_user.user_id)
        
        if not (hasattr(request.state, 'session') and request.state.session):
            # 当非session认证时，还需保存session，保证app中的代码的兼容性
            if self.current_user.jwt:
                request.state.session_sid = self.current_user.jwt
            # 临时session 不需要送cookie，避免浪费流量
            # 这里我们暂时保留会话数据，但不会自动发送cookie
            mysession = self.session_manager.get_session(self.current_user.jwt)
            
            # 如果只通过了Kong 认证，且没有在redis中存该session key ，如：key auth，则写入kong 认证的信息给session
            if (request.headers.get('X-Consumer-Username') and
                request.headers.get('X-Consumer-Custom-Id') and
                not mysession.get('_data') or len(mysession.get('_data', {})) == 0):
                # 有 X-Consumer-Username 可能是JWT和key auth, 没有mysession 只有key auth 认证时，才需要更新session
                if request.query_params.get('apikey') == self.current_user.jwt:
                    session_data = {
                        'uid': self.current_user.user_id,
                        'uname': self.current_user.user_name,
                        'systemuser': False,
                        'manageuser': False,
                        'manageuserid': '',
                        'type': 'appuser'
                    }
                    if hasattr(request.state, 'session'):
                        request.state.session.update(session_data)
                        request.state.session_modified = True
                else:  # logout 后redis session已经被清除, 但jwt 还在, 通过jwt传递参数的api,需要重新认证
                    self.current_user.user_name = None
                    self.current_user.user_id = None
                    return
            else:
                # 只有session 认证才需要设定这个redis的有效期
                # request.state.session['temp_expiration_time'] = session_expiration_time
                if hasattr(request.state, 'session') and mysession.get('_data'):
                    request.state.session.update(mysession.get('_data', {}))
                    request.state.session_modified = True
                    # 更新用户信息
                    self.current_user.user_name = mysession.get('_data', {}).get('uname', self.current_user.user_name)
                    self.current_user.user_id = mysession.get('_data', {}).get('uid', self.current_user.user_id)
        
        # 构建当前用户对象
        if hasattr(request.state, 'session') and request.state.session:
            self.current_user.current_user = User(
                uid=request.state.session.get('uid'),
                uname=request.state.session.get('uname'),
                systemuser=request.state.session.get('systemuser', False),
                manageuser=request.state.session.get('manageuser', False),
                manageuserid=request.state.session.get('manageuserid'),
                companyid=request.state.session.get('companyid', ''),
                type=request.state.session.get('type', 'appuser'),
                phone=request.state.session.get('phone', '')
            )
        else:
            # 使用从header或jwt获取的信息构建用户对象
            self.current_user.current_user = User(
                uid=self.current_user.user_id,
                uname=self.current_user.user_name,
                systemuser=False,
                manageuser=False,
                manageuserid='',
                companyid='',
                type='appuser',
                phone=''
            )
        if not self.current_user.user_name and not self.current_user.user_id:
            self.current_user.user_name = self.current_user.current_user.uname
            self.current_user.user_id = self.current_user.current_user.uid
        # 更新状态
        if hasattr(request.state, 'session') and request.state.session:
            request.state.session['uname'] = self.current_user.current_user.uname
            request.state.session['uid'] = self.current_user.current_user.uid
            request.state.session_modified = True

    async def valid_login(self, request: Request):
        """
        验证登录装饰器对应的依赖函数
        """
        self._set_auth(request)
        
        if self.current_user.user_name or self.current_user.user_id:
            # 将当前用户信息附加到请求状态中
            request.state.current_user = self.current_user.current_user
            request.state.user_name = self.current_user.user_name
            request.state.user_id = self.current_user.user_id
            request.state.jwt = self.current_user.jwt
            return self.current_user.current_user
        
        # 没有认证则抛出401异常
        raise HTTPException(status_code=401, detail="Unauthorized")

    def get_current_user(self):
        """
        获取当前用户的依赖函数
        """
        async def _get_current_user(request: Request):
            # 首先尝试获取当前用户，如果没有则先执行认证
            current_user = getattr(request.state, 'current_user', None)
            if not current_user:
                # 尝试通过认证获取用户
                await self.valid_login(request)
                current_user = getattr(request.state, 'current_user', None)
            
            if not current_user:
                raise HTTPException(status_code=401, detail="User not authenticated")
            return current_user
        return _get_current_user

    async def valid_sign(self, request: Request):
        """
        验证签名的依赖函数
        """
        # key认证的token为key，浏览器端因访问不到jwt，可以传username
        # json.dumps(json_body)+token(或用户名)+noncestr+timestamp 转md5
        sign = request.headers.get('X-MW-Sign') or request.query_params.get("sign")
        
        if not sign:
            raise HTTPException(status_code=400, detail="Sign parameter is required")
        
        # 获取当前用户信息
        current_user_obj = getattr(request.state, 'current_user', None)
        current_jwt = getattr(request.state, 'jwt', None)
        
        if not current_user_obj and not current_jwt:
            raise HTTPException(status_code=401, detail="Authentication required for sign validation")
        
        # 构建签名源字符串 - 使用用户名
        body_bytes = await request.body()
        sign_source_str = (
            f'{body_bytes.decode()}{current_user_obj.uname.lower() if current_user_obj else ""}'
            f'{request.headers.get("X-MW-Noncestr") or request.query_params.get("noncestr", "")}'
            f'{request.headers.get("X-MW-Timestamp") or request.query_params.get("timestamp", "")}'
        )
        sign_md5 = hashlib.md5(sign_source_str.encode()).hexdigest()
        
        print(f'sign_source_user:{sign_source_str}, sign_md5:{sign_md5}, sign_web:{sign}')
        
        if sign_md5 == sign:
            return True
        
        # 如果使用用户名验证失败，尝试使用JWT验证
        if current_jwt:
            sign_source_str_jwt = (
                f'{body_bytes.decode()}{current_jwt}'
                f'{request.headers.get("X-MW-Noncestr") or request.query_params.get("noncestr", "")}'
                f'{request.headers.get("X-MW-Timestamp") or request.query_params.get("timestamp", "")}'
            )
            sign_md5_jwt = hashlib.md5(sign_source_str_jwt.encode()).hexdigest()
            print(f'sign_source_jwt:{sign_source_str_jwt}, sign_md5_jwt:{sign_md5_jwt},sign_web:{sign}')
            
            if sign_md5_jwt == sign:
                return True
        
        # 两者都不匹配，返回错误
        raise HTTPException(
            status_code=400,
            detail=(
                f'sign error,web sign:{sign},'
                f'srv sign_user {sign_md5},'
                f'srv sign_jwt {sign_md5_jwt},'
                f'noncestr:{request.headers.get("X-MW-Noncestr") or request.query_params.get("noncestr", "")},'
                f'timestamp:{request.headers.get("X-MW-Timestamp") or request.query_params.get("timestamp", "")}'
            )
        )

    def valid_phone_code(self, phone_code_arg_name: str):
        """
        检查验证码的依赖函数，需要传入验证码的参数名称
        """
        async def _valid_phone_code(request: Request):
            # 从请求体或查询参数获取验证码
            body_data = {}
            try:
                body_data = await request.json()
            except:
                pass
            
            # 优先从路径参数中获取（因为依赖注入无法直接访问路径参数）
            # 但如果提供了body参数名，也从请求体中查找
            phone_code = body_data.get(phone_code_arg_name)
            
            if not phone_code:
                # 如果没在请求体中找到，尝试从查询参数中获取
                phone_code = request.query_params.get(phone_code_arg_name)
            
            if not phone_code:
                raise HTTPException(status_code=400, detail=f"Phone code parameter '{phone_code_arg_name}' is required")
            
            # 获取当前用户电话
            current_user_obj = getattr(request.state, 'current_user', None)
            if not current_user_obj or not current_user_obj.phone:
                raise HTTPException(status_code=401, detail="Phone not found in user profile")
            
            phone = current_user_obj.phone
            valid_code, hits = self.session_manager.redis.hmget(
                f'session:phone.valid.code:{phone}', 
                ['valid_code', 'hits']
            )
            
            if valid_code:
                # 增加最大尝试次数计数
                self.session_manager.redis.hincrby(
                    f'session:phone.valid.code:{phone}', 
                    'max_hits', 
                    1
                )
            
            if valid_code and valid_code.decode() == phone_code:
                # 验证成功，删除验证码
                self.session_manager.redis.delete(f'session:phone.valid.code:{phone}')
                return True
            
            # 验证失败，减少尝试次数
            self.session_manager.redis.hincrby(f'session:phone.valid.code:{phone}', 'hits', -1)
            
            # 检查剩余尝试次数
            if hits and int(hits.decode()) <= 0:
                self.session_manager.redis.delete(f'session:phone.valid.code:{phone}')
            
            raise HTTPException(
                status_code=400, 
                detail="The phone valid code is error, please retry enter."
            )
        
        return _valid_phone_code


# 便捷的依赖注入函数
def get_valid_login(auth_instance: FastApiKongAuth):
    """
    创建认证依赖注入函数的便捷方法
    """
    async def auth_dependency(request: Request):
        return await auth_instance.valid_login(request)
    return auth_dependency