import logging
# os is no longer needed for token persistence
import json
import time
import requests
from functools import wraps
from flask import current_app, redirect, url_for, session, request
# Import the mongo instance from your extensions file
from .extensions import mongo

logger = logging.getLogger(__name__)


def make_zoho_api_request(method, url_path, params=None, json_data=None, files=None):
    """
    Makes a request to the Zoho Books API, handling authentication and token refresh.
    It now supports multipart/form-data for file uploads via the 'files' parameter.

    For JSON responses, it returns (data, None).
    For file/image downloads, it returns (raw_bytes, headers).
    """
    access_token = get_access_token()
    if not access_token:
        # This will be caught by the calling function's try-except block
        raise Exception("Zoho authentication failed: Could not retrieve a valid access token.")

    # Dynamically determine the base URL from the stored token information
    stored_token = _get_stored_token()
    if not stored_token or 'api_domain' not in stored_token:
        logger.error("Zoho 'api_domain' not found in stored token. The token is outdated or invalid. Please re-authenticate.")
        clear_zoho_token() # Force re-authentication
        raise Exception("Zoho configuration error: API domain missing. Please log in again.")

    api_domain = stored_token['api_domain']
    base_url = f"{api_domain}/books/v3"
    
    headers = {
        "Authorization": f"Zoho-oauthtoken {access_token}",
    }
    # For multipart/form-data (when 'files' is used), 'requests' sets the Content-Type automatically.
    # For JSON, we set it explicitly.
    if json_data:
        headers["Content-Type"] = "application/json"

    full_url = f"{base_url}{url_path}"

    logger.debug(f"Making Zoho API request: {method} {full_url}")
    # Pass the 'files' parameter to the requests call for handling file uploads
    response = requests.request(method, full_url, headers=headers, params=params, json=json_data, files=files)

    # Check for success (200 OK, 201 Created)
    if response.status_code not in [200, 201]:
        error_message = f"Zoho API Error ({response.status_code}): {response.text}"
        logger.error(error_message)
        response.raise_for_status()
    
    content_type = response.headers.get('content-type', '').lower()
    # The print statement below was already present, I've left it as is.
    print(f"Response Content-Type: {content_type}")

    # If the response is not JSON, treat it as a file/image download.
    if 'application/json' not in content_type:
        logger.info(f"Received non-JSON content-type '{content_type}', returning raw content and headers.")
        return response.content, response.headers

    # Handle JSON responses, returning a tuple for consistency.
    if response.text:
        return response.json(), None

    # Handle empty success responses
    return None, None


def initialize_zoho_sdk(grant_token=None, accounts_server_url=None):
    """
    Initializes Zoho authentication. If a grant_token is provided, it exchanges
    it for access/refresh tokens using the correct data center.
    """
    if grant_token:
        logger.info("Initializing Zoho Auth with new GRANT token.")
        try:
            config = current_app.config
            
            # Use the dynamically provided accounts_server_url or default to .in
            base_accounts_url = accounts_server_url if accounts_server_url else 'https://accounts.zohocloud.ca'
            token_url = f'{base_accounts_url}/oauth/v2/token'
            logger.info(f"Using token exchange URL: {token_url}")

            payload = {
                'code': grant_token,
                'client_id': config['ZOHO_CLIENT_ID'],
                'client_secret': config['ZOHO_CLIENT_SECRET'],
                'redirect_uri': config['ZOHO_REDIRECT_URL'],
                'grant_type': 'authorization_code'
            }
            response = requests.post(token_url, data=payload)
            response.raise_for_status()
            token_data = response.json()
            
            # --- EDIT: Print the received datacenter URL and store it ---
            api_domain = token_data.get("api_domain")
            if api_domain:
                logger.info(f"Received Zoho API Domain: {api_domain}")
                print(f"--- Successfully received Zoho API Domain: {api_domain} ---")
            else:
                logger.warning("Could not find 'api_domain' in the token response. API requests might fail.")
                print("--- WARNING: 'api_domain' not found in Zoho's response. ---")

            # Store the accounts server URL for refreshing the token later
            token_data['accounts_server'] = base_accounts_url

            # Log and print refresh token if returned
            if 'refresh_token' in token_data:
                logger.info("Received refresh_token from Zoho during token exchange.")
                logger.info(f"refresh_token: {token_data.get('refresh_token')}")
                print(f"--- Received Zoho refresh_token: {token_data.get('refresh_token')} ---")
            else:
                logger.info("No refresh_token present in Zoho response for this exchange.")
                print("--- No refresh_token received from Zoho in this exchange. ---")

            logger.info(f"Token data received: {token_data}")
            token_data['expires_at'] = int(time.time()) + token_data['expires_in']
            _save_token(token_data)
            logger.info("Successfully exchanged grant token and stored tokens.")
        except requests.exceptions.RequestException as e:
            error_text = e.response.text if e.response else "No response from server"
            logger.critical(f"Failed to exchange grant token: {error_text}", exc_info=True)
            raise Exception(f"Failed to get token from Zoho: {error_text}") from e
    else:
        logger.debug("initialize_zoho_sdk called without grant token. No action taken.")

def _get_stored_token():
    """Reads token data from MongoDB."""
    try:
        # Use a consistent identifier for the single token document
        token_doc = mongo.db.zoho_tokens.find_one({'_id': 'zoho_oauth_token'})
        if token_doc:
            # The _id is not part of the token data itself, remove it
            token_doc.pop('_id', None)
        return token_doc
    except Exception as e:
        logger.error(f"Error reading Zoho token from MongoDB: {e}", exc_info=True)
        return None

def _save_token(token_data):
    """Saves token data to MongoDB."""
    try:
        # Preserve an existing refresh_token if the current token_data doesn't contain one.
        existing = mongo.db.zoho_tokens.find_one({'_id': 'zoho_oauth_token'})
        if existing and 'refresh_token' in existing and 'refresh_token' not in token_data:
            # Preserve the previously stored refresh token
            token_data['refresh_token'] = existing['refresh_token']
            logger.info("Preserved existing refresh_token because the new token response did not include one.")
            print("--- Preserved existing Zoho refresh_token (not returned by current response). ---")

        # Set expires_at if present as int already; otherwise leave as-is
        # Ensure the DB document uses fixed _id
        mongo.db.zoho_tokens.update_one(
            {'_id': 'zoho_oauth_token'},
            {'$set': token_data},
            upsert=True
        )
        logger.info("Zoho token saved to MongoDB (upsert).")
    except Exception as e:
        logger.error(f"Error saving Zoho token to MongoDB: {e}", exc_info=True)

def _refresh_access_token():
    """Uses the refresh token to get a new access token from the correct data center."""
    stored_token = _get_stored_token()
    if not stored_token or 'refresh_token' not in stored_token:
        logger.error("Refresh token not found. Cannot refresh.")
        return None

    logger.info("Refreshing Zoho access token.")
    try:
        config = current_app.config
        
        # Use the stored accounts_server or default for backward compatibility
        base_accounts_url = stored_token.get('accounts_server', 'https://accounts.zohocloud.ca')
        token_url = f'{base_accounts_url}/oauth/v2/token'
        logger.info(f"Using token refresh URL: {token_url}")

        payload = {
            'refresh_token': stored_token['refresh_token'],
            'client_id': config['ZOHO_CLIENT_ID'],
            'client_secret': config['ZOHO_CLIENT_SECRET'],
            'grant_type': 'refresh_token'
        }
        response = requests.post(token_url, data=payload)
        response.raise_for_status()
        new_token_data = response.json()

        # Preserve crucial details not always present in the refresh response
        new_token_data['refresh_token'] = new_token_data.get('refresh_token', stored_token['refresh_token'])
        new_token_data['accounts_server'] = new_token_data.get('accounts_server', base_accounts_url)
        # Use .get for expires_in to avoid KeyError if absent
        new_token_data['expires_at'] = int(time.time()) + int(new_token_data.get('expires_in', 3600))

        # Log whether Zoho returned a refresh_token on refresh calls (often not returned).
        if 'refresh_token' in new_token_data and new_token_data.get('refresh_token') != stored_token.get('refresh_token'):
            logger.info("Zoho returned a new refresh_token during refresh.")
            print(f"--- Zoho returned a new refresh_token during refresh: {new_token_data.get('refresh_token')} ---")
        else:
            logger.info("Zoho did not return a new refresh_token during refresh; preserving existing one.")
            print("--- Zoho did not return a new refresh_token during refresh; using stored refresh_token. ---")

        _save_token(new_token_data)
        logger.info("Successfully refreshed and saved new access token.")
        return new_token_data
    except requests.exceptions.RequestException as e:
        error_text = e.response.text if e.response else "No response from server"
        logger.critical(f"Failed to refresh Zoho access token: {error_text}", exc_info=True)
        clear_zoho_token()
        return None

def get_access_token():
    """Returns a valid access token, refreshing it if necessary."""
    token = _get_stored_token()
    if not token:
        return None

    if token.get('expires_at', 0) < time.time() + 60:
        logger.info("TRYING TO GENERATE A REFRESH TOKEN")
        token = _refresh_access_token()
        if not token:
            return None

    return token.get('access_token')

def is_zoho_token_available():
    """Checks if a refresh token has been persisted in MongoDB."""
    token = _get_stored_token()
    return token is not None and 'refresh_token' in token

def clear_zoho_token():
    """Removes the persisted token from MongoDB to effectively log out."""
    try:
        result = mongo.db.zoho_tokens.delete_one({'_id': 'zoho_oauth_token'})
        if result.deleted_count > 0:
            logger.info("Zoho token removed from MongoDB successfully.")
        else:
            logger.warning("Attempted to clear Zoho token, but none was found in MongoDB.")
    except Exception as e:
        logger.error(f"Error removing Zoho token from MongoDB: {e}", exc_info=True)


def zoho_token_required(function):
    """
    Decorator to ensure a valid Zoho token exists. If not, redirects to the login page.
    """
    @wraps(function)
    def decorator(*args, **kwargs):
        if not is_zoho_token_available():
            session['next_url'] = request.url
            return redirect(url_for("zoho.login"))
        return function(*args, **kwargs)
    return decorator