from flask import request, g
import datetime
import json
import logger_module.package_logger

# Save original print_log
original_print_log = logger_module.package_logger.print_log

def custom_print_log(message):
    try:
        from flask import has_request_context
        if has_request_context():
            if getattr(g, 'log_stream_name', 'mconnect') == 'lender':
                if isinstance(message, dict) and message.get('API_TRANSACTION') or (isinstance(message, str) and '"API_request":' in message):
                    pass # Keep the transaction JSON pure
                else:
                    endpoint = request.path.split('/')[-1] if request.path else 'unknown'
                    if isinstance(message, dict):
                        message = f"\n{json.dumps(message, indent=4)}\n - {endpoint}"
                    else:
                        message = f"{message} - {endpoint}"
    except Exception:
        pass
    
    original_print_log(message)

logger_module.package_logger.print_log = custom_print_log

def setup_logging_middleware(app):
    @app.before_request
    def before_request_handler():
        host_name = request.headers.get("Host", "").split(".")[0]
        # Fallback for local testing if needed
        if host_name in ['localhost', '127', '']:
            host_name = request.headers.get("X-Custom-Host", 'lending-uat')
            
        if 'lending' in host_name:
            g.log_stream_name = 'lender'
            g.request_start_time = datetime.datetime.now().isoformat()
            try:
                g.req_payload = request.get_json(silent=True, force=True) or {}
            except Exception:
                g.req_payload = {}
        else:
            g.log_stream_name = 'mconnect'

    @app.after_request
    def after_request_handler(response):
        if getattr(g, 'log_stream_name', '') == 'lender':
            try:
                # Retrieve the request payload safely captured earlier
                req_payload = getattr(g, 'req_payload', {})
                    
                # Safely get response JSON
                try:
                    resp_payload = json.loads(response.get_data(as_text=True))
                except Exception:
                    resp_payload = {}
                    
                # Extract application number
                app_no = req_payload.get('application_number') or req_payload.get('applicationNo')
                if not app_no and isinstance(resp_payload, dict):
                    app_no = resp_payload.get('application_number') or resp_payload.get('applicationNo')
                    
                # Inject application_number into request payload
                if app_no and isinstance(req_payload, dict):
                    req_payload['application_number'] = app_no
                    
                request_time = getattr(g, 'request_start_time', datetime.datetime.now().isoformat())
                response_time = datetime.datetime.now().isoformat()
                    
                transaction = {
                    "API_request": {
                        "api": request.path,
                        "status_code": response.status_code,
                        "application_number": app_no,
                        "request": req_payload,
                        "request_time": request_time,
                        "response": resp_payload,
                        "response_time": response_time
                    }
                }
                
                logger_module.package_logger.print_log(json.dumps(transaction))
            except Exception as e:
                logger_module.package_logger.print_log(f"Error logging API_request: {e}")
                
        return response