"""
Service for managing multiple admin NotebookLM accounts.
Provides random selection logic for load balancing.
"""
import os
import random
import json
from datetime import datetime, timedelta
from flask import current_app
from app.extensions import db
from app.models.nlm_admin_account import NLMAdminAccount


def get_admin_storage_path(account_id):
    """Get storage path for an admin account."""
    account_dir = f"/data/nlm_admin_{account_id}"
    os.makedirs(account_dir, exist_ok=True)
    return os.path.join(account_dir, "storage_state.json")


def add_admin_account(name, google_email, session_path=None, created_by_id=None):
    """
    Add a new admin NotebookLM account.

    Args:
        name: Display name for the account
        google_email: Google account email
        session_path: Path to storage_state.json (optional)
        created_by_id: ID of admin user who added this account

    Returns:
        NLMAdminAccount: Created account or None if failed
    """
    try:
        account = NLMAdminAccount(
            name=name,
            google_email=google_email,
            session_path=session_path,
            created_by_id=created_by_id,
            is_active=True,
            is_healthy=True,
        )
        db.session.add(account)
        db.session.commit()

        current_app.logger.info(f"Added admin NotebookLM account: {name} ({google_email})")
        return account

    except Exception as e:
        current_app.logger.error(f"Error adding admin account: {e}")
        db.session.rollback()
        return None


def get_available_admin_account():
    """
    Get a random available admin account for notebook creation.
    Uses weighted random selection favoring accounts with less usage.

    Returns:
        NLMAdminAccount: Available account or None if none available
    """
    try:
        # Get all active and healthy accounts
        accounts = NLMAdminAccount.query.filter_by(
            is_active=True,
            is_healthy=True
        ).all()

        # Filter out accounts at capacity
        available = [a for a in accounts if a.total_notebooks < 250]

        if not available:
            current_app.logger.warning("No available admin NotebookLM accounts (all at capacity)")
            return None

        # Weighted random selection - favor accounts with less usage
        # Weight is inversely proportional to notebook count
        weights = []
        for acc in available:
            # Base weight + bonus for low usage
            weight = max(1, 250 - acc.total_notebooks)
            # Additional weight based on priority (lower priority number = higher weight)
            priority_weight = max(1, 200 - acc.priority)
            weights.append(weight * priority_weight)

        # Normalize weights
        total_weight = sum(weights)
        if total_weight == 0:
            return random.choice(available)

        normalized_weights = [w / total_weight for w in weights]

        # Select account
        selected_index = random.choices(range(len(available)), weights=normalized_weights)[0]
        selected = available[selected_index]

        current_app.logger.info(f"Selected admin account: {selected.name} ({selected.total_notebooks}/250 notebooks)")
        return selected

    except Exception as e:
        current_app.logger.error(f"Error getting available admin account: {e}")
        return None


def mark_account_used(account_id):
    """Increment notebook count and update last_used_at for an account."""
    try:
        account = NLMAdminAccount.query.get(account_id)
        if account:
            account.total_notebooks += 1
            account.last_used_at = datetime.utcnow()
            db.session.commit()

    except Exception as e:
        current_app.logger.error(f"Error marking account used: {e}")


def mark_account_unhealthy(account_id, error_message=None):
    """Mark an account as unhealthy (session expired or error)."""
    try:
        account = NLMAdminAccount.query.get(account_id)
        if account:
            account.is_healthy = False
            account.error_message = error_message
            account.last_health_check = datetime.utcnow()
            db.session.commit()

            current_app.logger.warning(f"Marked account {account.name} as unhealthy: {error_message}")

    except Exception as e:
        current_app.logger.error(f"Error marking account unhealthy: {e}")


def mark_account_healthy(account_id):
    """Mark an account as healthy (session valid)."""
    try:
        account = NLMAdminAccount.query.get(account_id)
        if account:
            account.is_healthy = True
            account.error_message = None
            account.last_health_check = datetime.utcnow()
            db.session.commit()

    except Exception as e:
        current_app.logger.error(f"Error marking account healthy: {e}")


def get_all_admin_accounts():
    """Get all admin accounts."""
    return NLMAdminAccount.query.order_by(NLMAdminAccount.priority, NLMAdminAccount.created_at).all()


def get_admin_account(account_id):
    """Get a specific admin account."""
    return NLMAdminAccount.query.get(account_id)


def delete_admin_account(account_id):
    """Delete an admin account."""
    try:
        account = NLMAdminAccount.query.get(account_id)
        if account:
            # Delete session file if exists
            if account.session_path and os.path.exists(account.session_path):
                try:
                    os.remove(account.session_path)
                except:
                    pass

            db.session.delete(account)
            db.session.commit()

            current_app.logger.info(f"Deleted admin account: {account.name}")
            return True

        return False

    except Exception as e:
        current_app.logger.error(f"Error deleting admin account: {e}")
        db.session.rollback()
        return False


def update_admin_account(account_id, **kwargs):
    """Update admin account fields."""
    try:
        account = NLMAdminAccount.query.get(account_id)
        if account:
            for key, value in kwargs.items():
                if hasattr(account, key):
                    setattr(account, key, value)

            account.updated_at = datetime.utcnow()
            db.session.commit()

            return account

        return None

    except Exception as e:
        current_app.logger.error(f"Error updating admin account: {e}")
        db.session.rollback()
        return None


def check_account_health(account_id):
    """
    Check if an account's session is still valid.

    Args:
        account_id: ID of the account to check

    Returns:
        bool: True if session is valid, False otherwise
    """
    try:
        account = NLMAdminAccount.query.get(account_id)
        if not account or not account.session_path:
            return False

        # Check if session file exists and has SID cookie
        if not os.path.exists(account.session_path):
            return False

        with open(account.session_path, 'r') as f:
            storage = json.load(f)

        cookies = storage.get('cookies', [])
        cookie_names = [c.get('name') for c in cookies]

        has_sid = 'SID' in cookie_names or 'HSID' in cookie_names or 'SSID' in cookie_names

        if has_sid:
            mark_account_healthy(account_id)
            return True
        else:
            mark_account_unhealthy(account_id, "Missing SID cookies")
            return False

    except Exception as e:
        current_app.logger.error(f"Error checking account health: {e}")
        mark_account_unhealthy(account_id, str(e))
        return False


def check_all_accounts_health():
    """Check health of all admin accounts."""
    accounts = get_all_admin_accounts()
    results = {}

    for account in accounts:
        results[account.id] = check_account_health(account.id)

    return results


def get_account_usage_stats():
    """Get usage statistics for all admin accounts."""
    accounts = get_all_admin_accounts()

    total_notebooks = sum(a.total_notebooks for a in accounts)
    total_capacity = len(accounts) * 250
    usage_percentage = int((total_notebooks / total_capacity * 100)) if total_capacity > 0 else 0

    return {
        "total_accounts": len(accounts),
        "active_accounts": sum(1 for a in accounts if a.is_active),
        "healthy_accounts": sum(1 for a in accounts if a.is_healthy),
        "available_accounts": sum(1 for a in accounts if a.is_available),
        "total_notebooks": total_notebooks,
        "total_capacity": total_capacity,
        "usage_percentage": usage_percentage,
        "accounts": [a.to_dict() for a in accounts],
    }


def get_least_used_account():
    """Get the admin account with the least number of notebooks."""
    return NLMAdminAccount.query.filter_by(
        is_active=True,
        is_healthy=True
    ).order_by(NLMAdminAccount.total_notebooks.asc()).first()


def authenticate_admin_account_browser(account_id):
    """
    Trigger browser authentication for an admin account.
    This is used when the session expires and needs to be refreshed.

    Args:
        account_id: ID of the account to authenticate

    Returns:
        bool: True if authentication successful, False otherwise
    """
    try:
        from app.ai.nlm_undetected import authenticate_user_notebooklm

        account = NLMAdminAccount.query.get(account_id)
        if not account:
            return False

        storage_path = get_admin_storage_path(account_id)

        # Run authentication
        success = authenticate_user_notebooklm(storage_path)

        if success:
            # Update account with new session
            account.session_path = storage_path
            account.is_healthy = True
            account.error_message = None
            account.last_health_check = datetime.utcnow()
            db.session.commit()

            current_app.logger.info(f"Successfully authenticated admin account: {account.name}")
            return True
        else:
            mark_account_unhealthy(account_id, "Authentication failed")
            return False

    except Exception as e:
        current_app.logger.error(f"Error authenticating admin account: {e}")
        mark_account_unhealthy(account_id, str(e))
        return False