import json
import time
from huggingface_hub import HfApi, hf_hub_download
from werkzeug.security import generate_password_hash, check_password_hash
import gradio as gr
from retry import retry  # Requires `pip install retry`

# Handle HfHubHTTPError import gracefully (it's in utils, and may vary by version)
try:
    from huggingface_hub.utils import HfHubHTTPError
except ImportError:
    from requests.exceptions import HTTPError as HfHubHTTPError  # Fallback to base HTTPError for compatibility

# ================== ALLOWED USERS + AUDIT LOG ==================
def is_allowed_username(username, auth_globals):
    username = username.strip().lower()
    try:
        path = hf_hub_download(repo_id=auth_globals["user_repo"], filename="allowed_users.json", repo_type="dataset", token=auth_globals["hf_token"])
        with open(path, 'r') as f:
            allowed = json.load(f)
        return username in [u.strip().lower() for u in allowed]
    except HfHubHTTPError as e:
        gr.Warning(f"Failed to check allowed users: {str(e)}")
        return False
    except Exception as e:
        gr.Warning(f"Unexpected error in is_allowed_username: {str(e)}")
        return False

@retry(tries=3, delay=1, backoff=2)  # Retry on transient errors
def log_activity(action, username, details="", auth_globals=None):
    if not auth_globals or not username:
        return
    try:
        hf_api = HfApi(token=auth_globals["hf_token"])
        user_repo = auth_globals["user_repo"]
        try:
            path = hf_hub_download(repo_id=user_repo, filename="activity_log.json", repo_type="dataset", token=auth_globals["hf_token"])
            with open(path, 'r') as f:
                log = json.load(f)
        except HfHubHTTPError:
            log = []  # Initialize if file doesn't exist
        entry = {
            "timestamp": time.time(),
            "datetime": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
            "action": action,
            "username": username.strip().lower(),
            "details": details
        }
        log.append(entry)
        hf_api.upload_file(
            path_or_fileobj=json.dumps(log, indent=2).encode(),
            path_in_repo="activity_log.json",
            repo_id=user_repo,
            repo_type="dataset",
            token=auth_globals["hf_token"]
        )
    except Exception as e:
        gr.Warning(f"Failed to log activity: {str(e)}")  # Better than silent failure

def get_total_accounts(auth_globals):
    try:
        hf_api = HfApi(token=auth_globals["hf_token"])
        files = hf_api.list_repo_files(repo_id=auth_globals["user_repo"], repo_type="dataset")
        return sum(1 for f in files if f.endswith('.json') and f not in ["allowed_users.json", "activity_log.json"])
    except HfHubHTTPError as e:
        gr.Warning(f"Failed to get total accounts: {str(e)}")
        return 0
    except Exception as e:
        gr.Warning(f"Unexpected error in get_total_accounts: {str(e)}")
        return 0

# ================== HELPER FOR FAILURE RETURNS ==================
def failure_return(logged_in_username, message, user_chats=None, current_title=None, history=None, dropdown=None, context_tb=None, save_btn=None):
    # Defaults are set to None and computed here to avoid NameError on evaluation
    if user_chats is None:
        user_chats = []
    if current_title is None:
        current_title = None
    if history is None:
        history = []
    if dropdown is None:
        dropdown = gr.update(choices=[], value=None)
    if context_tb is None:
        context_tb = gr.update(visible=bool(logged_in_username), value="")
    if save_btn is None:
        save_btn = gr.update(visible=bool(logged_in_username))

    return (
        user_chats, current_title, history, dropdown, logged_in_username, f"❌ {message}",
        gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False), context_tb, save_btn
    )

# ================== REGISTER ==================
def register_user(username, password, logged_in_username, auth_globals):
    hf_token = auth_globals["hf_token"]
    user_repo = auth_globals["user_repo"]
    hf_api = HfApi(token=hf_token)

    username = username.strip().lower()
    if not username.isalnum():  # Added: Sanitize username (alphanumeric only)
        return failure_return(logged_in_username, "Username must be alphanumeric.")

    if get_total_accounts(auth_globals) >= 20:
        return failure_return(logged_in_username, "Maximum 20 accounts reached.")

    if not username or not password:
        return failure_return(logged_in_username, "Please fill username and password.")

    if not is_allowed_username(username, auth_globals):
        return failure_return(logged_in_username, "This username is not pre-approved by admin.")

    try:
        # Check if already exists
        hf_hub_download(repo_id=user_repo, filename=f"{username}.json", repo_type="dataset", token=hf_token)
        return failure_return(logged_in_username, "User already exists!")
    except HfHubHTTPError:  # Expected if file doesn't exist
        pass
    except Exception as e:
        return failure_return(logged_in_username, f"Registration check failed: {str(e)}")

    try:
        user_data = {
            "password": generate_password_hash(password),
            "chats": []
        }
        hf_api.upload_file(
            path_or_fileobj=json.dumps(user_data).encode('utf-8'),
            path_in_repo=f"{username}.json",
            repo_id=user_repo,
            repo_type="dataset",
            token=hf_token
        )

        log_activity("REGISTER", username, "New account created", auth_globals)

        return [], None, [], gr.update(choices=[], value=None), username, "✅ Registration successful! Logged in.", gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True, value=""), gr.update(visible=True)
    except Exception as e:
        return failure_return(logged_in_username, f"Registration failed: {str(e)}")

# ================== LOGIN ==================
def login_user(username, password, logged_in_username, auth_globals):
    hf_token = auth_globals["hf_token"]
    user_repo = auth_globals["user_repo"]
    hf_api = HfApi(token=hf_token)

    username = username.strip().lower()
    if not username or not password:
        return failure_return(logged_in_username, "Please enter username and password.")

    if not is_allowed_username(username, auth_globals):
        return failure_return(logged_in_username, "This username is not pre-approved by admin.")

    try:
        user_file = hf_hub_download(repo_id=user_repo, filename=f"{username}.json", repo_type="dataset", token=hf_token)
        with open(user_file, 'r') as f:
            data = json.load(f)

        if check_password_hash(data["password"], password):
            chats = data.get("chats", [])
            if chats:
                current_title = chats[-1]["title"]
                history = chats[-1]["history"]
                context = chats[-1].get("context", "")  # Added: Load context
                choices = [chat["title"] for chat in chats]
                value = current_title
            else:
                current_title = None
                history = []
                context = ""
                choices = []
                value = None

            dropdown = gr.update(choices=choices, value=value)
            log_activity("LOGIN", username, "Successful login", auth_globals)
            return chats, current_title, history, dropdown, username, "✅ Login successful!", gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True, value=context), gr.update(visible=True)
        else:
            return failure_return(logged_in_username, "Incorrect password.")
    except HfHubHTTPError as e:
        return failure_return(logged_in_username, f"Login failed: User not found ({str(e)}).")
    except Exception as e:
        return failure_return(logged_in_username, f"Login failed: {str(e)}")

# ================== LOGOUT ==================
def logout_user():
    return [], None, [], gr.update(choices=[], value=None), None, "Logged out successfully!", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False, value=""), gr.update(visible=False)

# ================== SAVE CHATS ==================
@retry(tries=3, delay=1, backoff=2)  # Added retry for reliability
def save_chats(username, chats, auth_globals):
    if not username:
        return "No user logged in!"
    try:
        hf_token = auth_globals["hf_token"]
        user_repo = auth_globals["user_repo"]
        hf_api = HfApi(token=hf_token)
        user_file = hf_hub_download(repo_id=user_repo, filename=f"{username}.json", repo_type="dataset", token=hf_token)
        with open(user_file, 'r') as f:
            data = json.load(f)
        data["chats"] = chats  # Chats now include "context"
        hf_api.upload_file(path_or_fileobj=json.dumps(data).encode('utf-8'), path_in_repo=f"{username}.json", repo_id=user_repo, repo_type="dataset", token=hf_token)
        log_activity("SAVE_CHATS", username, "Chats updated", auth_globals)  # Added logging for saves
        return "Chats saved automatically!"
    except HfHubHTTPError as e:
        return f"Failed to save chats: {str(e)} (HF API error)"
    except Exception as e:
        return f"Failed to save chats: {str(e)}"

# ================== DELETE ACCOUNT ==================
def delete_account(confirm_username, logged_in_username, auth_globals):
    confirm_username = confirm_username.strip().lower()
    if not logged_in_username or confirm_username != logged_in_username:
        return "❌ Username does not match.", logged_in_username, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), None, None, None, None, gr.update(visible=True), gr.update(visible=True)

    if not is_allowed_username(logged_in_username, auth_globals):
        return "❌ Not allowed.", logged_in_username, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), None, None, None, None, gr.update(visible=True), gr.update(visible=True)

    hf_token = auth_globals["hf_token"]
    user_repo = auth_globals["user_repo"]
    hf_api = HfApi(token=hf_token)
    try:
        hf_api.delete_file(path_in_repo=f"{logged_in_username}.json", repo_id=user_repo, repo_type="dataset", token=hf_token)
        log_activity("DELETE_ACCOUNT", logged_in_username, "Account permanently deleted", auth_globals)
        return "✅ Account permanently deleted.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), [], None, [], gr.update(choices=[], value=None), gr.update(visible=False, value=""), gr.update(visible=False)
    except HfHubHTTPError as e:
        return f"❌ Deletion failed: {str(e)}", logged_in_username, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), None, None, None, None, gr.update(visible=True), gr.update(visible=True)
    except Exception as e:
        return f"❌ Deletion failed: {str(e)}", logged_in_username, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), None, None, None, None, gr.update(visible=True), gr.update(visible=True)