import inspect
from flask import Blueprint, request, jsonify, current_app, url_for, redirect, g
import jwt
from datetime import datetime, timedelta
from .db import Database
from .models import User, Role, ApiToken
from .exceptions import AuthError
from .decorators import public_endpoint
import uuid
import requests
import bcrypt
import logging
import os
from functools import wraps

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

def handle_auth_errors(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        try:
            return f(*args, **kwargs)
        except AuthError as e:
            response = jsonify(e.to_dict())
            response.status_code = e.status_code
            return response
    return decorated

class AuthManager:
    def __init__(self, app=None, db_dsn=None, jwt_secret=None, oauth_config=None, id_type='integer'):
        logger.info("INITIALIZING AUTHMANAGER {} - {} - {}".format(db_dsn, jwt_secret, not app))
        self.db = Database(db_dsn, id_type=id_type) if db_dsn else None
        self.jwt_secret = jwt_secret
        logger.debug(f"Initializing AuthManager with JWT secret: {jwt_secret[:5]}..." if jwt_secret else "No JWT secret provided")
        self.oauth_config = oauth_config or {}
        
        if app:
            self.init_app(app)

    def _extract_token_from_header(self):
        auth = request.authorization
        if not auth or not auth.token:
            raise AuthError('No authorization header or token', 401)

        if auth.type.lower() != 'bearer':
            raise AuthError('Invalid authorization scheme', 401)

        return auth.token

    def get_redirect_uri(self):
        redirect_uri = os.getenv('REDIRECT_URL') or url_for('auth.oauth_callback', _external=True).replace("http://", "https://")
        logger.info(f"REDIRECT URI..: {redirect_uri}")
        return redirect_uri

    def _validate_api_token(self, api_token):
        try:
            parsed = ApiToken.parse_token(api_token)
            with self.db.get_cursor() as cur:
                # First get the API token record
                cur.execute("""
                    SELECT t.*, u.* FROM api_tokens t
                    JOIN users u ON t.user_id = u.id
                    WHERE t.id = %s
                """, (parsed['id'],))
                result = cur.fetchone()
                if not result:
                    raise AuthError('Invalid API token')

                # Verify the nonce
                if not bcrypt.checkpw(parsed['nonce'].encode('utf-8'), result['token'].encode('utf-8')):
                    raise AuthError('Invalid API token')

                # Check if token is expired
                if result['expires_at'] and result['expires_at'] < datetime.utcnow():
                    raise AuthError('API token has expired')

                # Update last used timestamp
                cur.execute("""
                    UPDATE api_tokens 
                    SET last_used_at = %s
                    WHERE id = %s
                """, (datetime.utcnow(), parsed['id']))

                # Fetch roles
                cur.execute("""
                    SELECT r.name FROM roles r
                    JOIN user_roles ur ON ur.role_id = r.id
                    WHERE ur.user_id = %s
                """, (result['user_id'],))
                roles = [row['name'] for row in cur.fetchall()]

                # Construct user object
                return {
                    'id': result['user_id'],
                    'username': result['username'],
                    'email': result['email'],
                    'real_name': result['real_name'],
                    'roles': roles
                }
        except ValueError:
            raise AuthError('Invalid token format')

    def _authenticate_request(self):
        auth_header = request.headers.get('Authorization')
        api_token = request.headers.get('X-API-Token')

        if auth_header and auth_header.startswith('Bearer '):
            # JWT authentication
            token = self._extract_token_from_header()
            return self.validate_token(token)
        elif api_token:
            # API token authentication
            return self._validate_api_token(api_token)
        else:
            raise AuthError('No authentication provided', 401)

    def require_auth(self, f):
        @wraps(f)
        def decorated(*args, **kwargs):
            user = self._authenticate_request()
            sig = inspect.signature(f)
            if 'requesting_user' in sig.parameters:
                kwargs['requesting_user'] = user

            return f(*args, **kwargs)
        return decorated
    
    def init_app(self, app):
        app.auth_manager = self
        app.register_blueprint(self.create_blueprint())

    def create_blueprint(self):
        bp = Blueprint('auth', __name__, url_prefix='/api/v1/users')

        @bp.before_request
        def load_user():
            view = current_app.view_functions.get(request.endpoint)
            if getattr(view, '_auth_public', False):
                return
            try:
                g.current_user = self._authenticate_request()
            except AuthError as e:
                response = jsonify(e.to_dict())
                response.status_code = e.status_code
                return response

        @bp.route('/login', methods=['POST'])
        @public_endpoint
        @handle_auth_errors
        def login():
            data = request.get_json()
            username = data.get('username')
            password = data.get('password')
            
            if not username or not password:
                raise AuthError('Username and password required', 400)
            
            with self.db.get_cursor() as cur:
                cur.execute("SELECT * FROM users WHERE username = %s", (username,))
                user = cur.fetchone()
                
                if not user or not self._verify_password(password, user['password_hash']):
                    raise AuthError('Invalid username or password', 401)
                
                # Fetch roles
                cur.execute("""
                    SELECT r.name FROM roles r
                    JOIN user_roles ur ON ur.role_id = r.id
                    WHERE ur.user_id = %s
                """, (user['id'],))
                roles = [row['name'] for row in cur.fetchall()]
                user['roles'] = roles
                
                token = self._create_token(user)
                refresh_token = self._create_refresh_token(user)
                
                return jsonify({
                    'token': token,
                    'refresh_token': refresh_token,
                    'user': user
                })

        @bp.route('/login/oauth', methods=['POST'])
        @public_endpoint
        @handle_auth_errors
        def oauth_login():
            provider = request.json.get('provider')
            if provider not in self.oauth_config:
                raise AuthError('Invalid OAuth provider', 400)

            redirect_uri = self.get_redirect_uri()
            return jsonify({
                'redirect_url': self._get_oauth_url(provider, redirect_uri)
            })

        @bp.route('/login/oauth2callback')
        @public_endpoint
        @handle_auth_errors
        def oauth_callback():
            code = request.args.get('code')
            provider = request.args.get('state')
            
            if not code or not provider:
                raise AuthError('Invalid OAuth callback', 400)

            user_info = self._get_oauth_user_info(provider, code)
            token = self._create_token(user_info)
            refresh_token = self._create_refresh_token(user_info)
            
            # Redirect to frontend with tokens
            frontend_url = os.getenv('FRONTEND_URL', 'http://localhost:5173')
            return redirect(f"{frontend_url}/oauth-callback?token={token}&refresh_token={refresh_token}")

        @bp.route('/login/profile')
        @handle_auth_errors
        def profile():
            token = request.headers.get('Authorization', '').split(' ')[-1]
            user = self.validate_token(token)
            return jsonify(user)

        @bp.route('/api-tokens', methods=['GET'])
        @handle_auth_errors
        @self.require_auth
        def get_tokens(requesting_user):
            tokens = self.get_user_api_tokens(requesting_user['id'])
            return jsonify(tokens)

        @bp.route('/api-tokens', methods=['POST'])
        @handle_auth_errors
        @self.require_auth
        def create_token(requesting_user):
            name = request.json.get('name')
            expires_in_days = request.json.get('expires_in_days')
            if not name:
                raise AuthError('Token name is required', 400)
            api_token = self.create_api_token(requesting_user['id'], name, expires_in_days)
            return jsonify({
                'id': api_token.id,
                'name': api_token.name,
                'token': api_token.get_full_token(),
                'created_at': api_token.created_at,
                'expires_at': api_token.expires_at
            })

        @bp.route('/token-refresh', methods=['POST'])
        @public_endpoint
        @handle_auth_errors
        def refresh_token():
            refresh_token = request.json.get('refresh_token')
            if not refresh_token:
                raise AuthError('No refresh token provided', 400)

            try:
                payload = jwt.decode(refresh_token, self.jwt_secret, algorithms=['HS256'])
                user_id = payload['sub']
                
                with self.db.get_cursor() as cur:
                    cur.execute("SELECT * FROM users WHERE id = %s", (user_id,))
                    user = cur.fetchone()

                if not user:
                    raise AuthError('User not found', 404)

                return jsonify({
                    'token': self._create_token(user),
                    'refresh_token': self._create_refresh_token(user)
                })
            except jwt.InvalidTokenError:
                raise AuthError('Invalid refresh token', 401)

        @bp.route('/api-tokens', methods=['POST'])
        @handle_auth_errors
        @self.require_auth
        def create_api_token(requesting_user):
            name = request.json.get('name')
            if not name:
                raise AuthError('Token name required', 400)

            token = self.create_api_token(requesting_user['id'], name)
            return jsonify({'token': token.token})

        @bp.route('/api-tokens/validate', methods=['GET'])
        @handle_auth_errors
        @self.require_auth
        def validate_api_token(requesting_user):
            token = request.json.get('token')
            if not token:
                raise AuthError('No API token provided', 401)
            token = ApiToken.parse_token_id(token)

            with self.db.get_cursor() as cur:
                cur.execute("""
                    SELECT * FROM api_tokens 
                    WHERE user_id = %s AND id = %s
                """, (requesting_user['id'], token))
                api_token = cur.fetchone()

            if not api_token:
                raise AuthError('Invalid API token', 401)

            # Check if token is expired
            if api_token['expires_at'] and api_token['expires_at'] < datetime.utcnow():
                raise AuthError('API token has expired', 401)

            # Update last used timestamp
            with self.db.get_cursor() as cur:
                cur.execute("""
                    UPDATE api_tokens 
                    SET last_used_at = %s
                    WHERE id = %s
                """, (datetime.utcnow(), api_token['id']))

            return jsonify({'valid': True})

        @bp.route('/api-tokens', methods=['DELETE'])
        @handle_auth_errors
        @self.require_auth
        def delete_api_token(requesting_user):
            token = request.json.get('token')
            if not token:
                raise AuthError('Token required', 400)
            token = ApiToken.parse_token_id(token)

            with self.db.get_cursor() as cur:
                cur.execute("""
                    DELETE FROM api_tokens 
                    WHERE user_id = %s AND id = %s
                    RETURNING id
                """, (requesting_user['id'], token))
                deleted_id = cur.fetchone()
                if not deleted_id:
                    raise ValueError('Token not found or already deleted')

            return jsonify({'deleted': True})

        @bp.route('/register', methods=['POST'])
        @public_endpoint
        @handle_auth_errors
        def register():
            data = request.get_json()
            
            # Hash the password
            password = data.get('password')
            if not password:
                raise AuthError('Password is required', 400)
            
            salt = bcrypt.gensalt()
            password_hash = bcrypt.hashpw(password.encode('utf-8'), salt)
            
            user = User(
                username=data['username'],
                email=data['email'],
                real_name=data['real_name'],
                roles=data.get('roles', []),
                id_generator=self.db.get_id_generator()
            )

            with self.db.get_cursor() as cur:
                if user.id is None:
                    cur.execute("""
                        INSERT INTO users (username, email, real_name, password_hash, created_at, updated_at)
                        VALUES (%s, %s, %s, %s, %s, %s)
                        RETURNING id
                    """, (user.username, user.email, user.real_name, password_hash.decode('utf-8'),
                          user.created_at, user.updated_at))
                    user.id = cur.fetchone()['id']
                else:
                    cur.execute("""
                        INSERT INTO users (id, username, email, real_name, password_hash, created_at, updated_at)
                        VALUES (%s, %s, %s, %s, %s, %s, %s)
                    """, (user.id, user.username, user.email, user.real_name, password_hash.decode('utf-8'),
                          user.created_at, user.updated_at))

            return jsonify({'id': user.id}), 201

        @bp.route('/roles', methods=['GET'])
        @public_endpoint
        @handle_auth_errors
        def get_roles():
            with self.db.get_cursor() as cur:
                cur.execute("SELECT * FROM roles")
                roles = cur.fetchall()
            return jsonify(roles)

        return bp

    def validate_token(self, token):
        try:
            logger.debug(f"Validating token: {token}")
            payload = jwt.decode(token, self.jwt_secret, algorithms=['HS256'])
            logger.debug(f"Token payload: {payload}")
            user_id = int(payload['sub'])  # Convert string ID back to integer
            
            with self.db.get_cursor() as cur:
                cur.execute("SELECT * FROM users WHERE id = %s", (user_id,))
                user = cur.fetchone()
                if not user:
                    logger.error(f"User not found for ID: {user_id}")
                    raise AuthError('User not found', 404)
                # Fetch roles
                cur.execute("""
                    SELECT r.name FROM roles r
                    JOIN user_roles ur ON ur.role_id = r.id
                    WHERE ur.user_id = %s
                """, (user_id,))
                roles = [row['name'] for row in cur.fetchall()]
                user['roles'] = roles

            return user
        except jwt.InvalidTokenError as e:
            logger.error(f"Invalid token error: {str(e)}")
            raise AuthError('Invalid token', 401)
        except Exception as e:
            logger.error(f"Unexpected error during token validation: {str(e)}")
            raise AuthError(str(e), 500)

    def get_current_user(self):
        if hasattr(g, 'current_user'):
            return g.current_user
        return self._authenticate_request()

    def get_user_api_tokens(self, user_id):
        """Get all API tokens for a user."""
        with self.db.get_cursor() as cur:
            cur.execute("""
                SELECT id, name, created_at, expires_at, last_used_at
                FROM api_tokens 
                WHERE user_id = %s
                ORDER BY created_at DESC
            """, (user_id,))
            return cur.fetchall()

    def create_api_token(self, user_id, name, expires_in_days=None):
        """Create a new API token for a user."""
        token = ApiToken(user_id, name, expires_in_days)
        
        with self.db.get_cursor() as cur:
            cur.execute("""
                INSERT INTO api_tokens (id, user_id, name, token, created_at, expires_at)
                VALUES (%s, %s, %s, %s, %s, %s)
            """, (token.id, token.user_id, token.name, token.token, token.created_at, token.expires_at))
            return token

    def _create_token(self, user):
        payload = {
            'sub': str(user['id']),
            'exp': datetime.utcnow() + timedelta(hours=1),
            'iat': datetime.utcnow()
        }
        logger.debug(f"Creating token with payload: {payload}")
        token = jwt.encode(payload, self.jwt_secret, algorithm='HS256')
        logger.debug(f"Created token: {token}")
        return token

    def _create_refresh_token(self, user):
        payload = {
            'sub': str(user['id']),
            'exp': datetime.utcnow() + timedelta(days=30),
            'iat': datetime.utcnow()
        }
        return jwt.encode(payload, self.jwt_secret, algorithm='HS256')

    def _verify_password(self, password, password_hash):
        return bcrypt.checkpw(password.encode('utf-8'), password_hash.encode('utf-8'))

    def _get_oauth_url(self, provider, redirect_uri):
        if provider == 'google':
            client_id = self.oauth_config['google']['client_id']
            scope = 'openid email profile'
            state = provider  # Pass provider as state for callback
            return f'https://accounts.google.com/o/oauth2/v2/auth?client_id={client_id}&redirect_uri={redirect_uri}&response_type=code&scope={scope}&state={state}'
        raise AuthError('Invalid OAuth provider')


    def _get_oauth_user_info(self, provider, code):
        if provider == 'google':
            client_id = self.oauth_config['google']['client_id']
            client_secret = self.oauth_config['google']['client_secret']
            redirect_uri = self.get_redirect_uri()

            # Exchange code for tokens
            token_url = 'https://oauth2.googleapis.com/token'
            token_data = {
                'client_id': client_id,
                'client_secret': client_secret,
                'code': code,
                'grant_type': 'authorization_code',
                'redirect_uri': redirect_uri
            }
            token_response = requests.post(token_url, data=token_data)
            logger.info("TOKEN RESPONSE: {} {} {} [[[{}]]]".format(token_response.text, token_response.status_code, token_response.headers, token_data))
            token_response.raise_for_status()
            tokens = token_response.json()

            # Get user info
            userinfo_url = 'https://www.googleapis.com/oauth2/v3/userinfo'
            userinfo_response = requests.get(
                userinfo_url,
                headers={'Authorization': f"Bearer {tokens['access_token']}"}
            )
            userinfo_response.raise_for_status()
            userinfo = userinfo_response.json()

            # Create or update user
            with self.db.get_cursor() as cur:
                cur.execute("SELECT * FROM users WHERE email = %s", (userinfo['email'],))
                user = cur.fetchone()

                if not user:
                    # Create new user
                    user = User(
                        username=userinfo['email'],
                        email=userinfo['email'],
                        real_name=userinfo.get('name', userinfo['email']),
                        id_generator=self.db.get_id_generator()
                    )
                    cur.execute("""
                        INSERT INTO users (username, email, real_name, created_at, updated_at)
                        VALUES (%s, %s, %s, %s, %s)
                        RETURNING id
                    """, (user.username, user.email, user.real_name, 
                          user.created_at, user.updated_at))
                    user.id = cur.fetchone()['id']
                    user = {'id': user.id, 'username': user.username, 'email': user.email, 
                           'real_name': user.real_name, 'roles': []}
                else:
                    # Update existing user
                    cur.execute("""
                        UPDATE users 
                        SET real_name = %s, updated_at = %s
                        WHERE email = %s
                    """, (userinfo.get('name', userinfo['email']), datetime.utcnow(), userinfo['email']))
                    user['real_name'] = userinfo.get('name', userinfo['email'])

            return user
        raise AuthError('Invalid OAuth provider') 
