from functools import wraps
from flask import g, session, redirect, url_for
from app import mysql

def company_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        company_id = session.get('company_id')
        if not company_id:
            return redirect(url_for('auth.login'))
        
        cursor = mysql.connection.cursor()
        cursor.execute('SELECT * FROM companies WHERE id = %s', (company_id,))
        company = cursor.fetchone()
        cursor.close()
        
        if not company:
            return redirect(url_for('auth.login'))
            
        g.company = {
            'id': company[0],
            'name': company[1],
            'email': company[2],
            'instructions': company[3],
            'website': company[4]
        }
        
        return f(*args, **kwargs)
    return decorated_function 