import os
import jwt
from functools import wraps
from flask import request, jsonify

# Should ideally be set in .env
SECRET_KEY = os.getenv('JWT_SECRET', 'super-secret-jwt-key-automarket')

def token_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = None
        auth_header = request.headers.get('Authorization')
        
        if auth_header and auth_header.startswith('Bearer '):
            token = auth_header.split(" ")[1]
        
        # Fallback to query param for file downloads
        if not token:
            token = request.args.get('token')
            
        if not token:
            return jsonify({'message': 'Token is missing!'}), 401

        try:
            data = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
            user_info = {
                'user_id': data['user_id'],
                'company_id': data['company_id'],
                'role': data['role']
            }
        except Exception as e:
            return jsonify({'message': 'Token is invalid or expired!'}), 401

        return f(user_info, *args, **kwargs)
    return decorated