import random
from funcs import *

the_reachable_void()

import asyncio
import json
import os
import uuid
import bcrypt
import time
import secrets
import hashlib
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import Response
from fastapi.middleware.cors import CORSMiddleware
import urllib.request

try:
    from huggingface_hub import HfApi, hf_hub_url, hf_hub_download
except ImportError:
    HfApi = None
    hf_hub_url = None
    hf_hub_download = None

# Environment variables configuration for Hugging Face storage
HF_TOKEN = os.getenv("database_token")
DATASET_REPO_ID = os.getenv("dataset_link")
USERS_FILENAME = "users.json"
CONVERSATIONS_FILENAME = "conversations.json"
SALTS_FILENAME = "remember_me_salts.json"
RECENTS_FILENAME = "users_recent_chats.json"
OFFLINE_QUEUES_FILENAME = "offline_queues.json"

os.makedirs("/data", exist_ok=True)
LOCAL_USERS_PATH = os.path.join("/data", USERS_FILENAME)
LOCAL_CONVERSATIONS_PATH = os.path.join("/data", CONVERSATIONS_FILENAME)
LOCAL_SALTS_PATH = os.path.join("/data", SALTS_FILENAME)
LOCAL_RECENTS_PATH = os.path.join("/data", RECENTS_FILENAME)
LOCAL_OFFLINE_QUEUES_PATH = os.path.join("/data", OFFLINE_QUEUES_FILENAME)

hf_api = HfApi(token=HF_TOKEN) if HF_TOKEN and HfApi else None

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

MAINTENANCE_FILENAME = "maintenance.json"
LOCAL_MAINTENANCE_PATH = os.path.join("/data", MAINTENANCE_FILENAME)

PENDING_USERS_SYNC = False
PENDING_CONVERSATIONS_SYNC = False
PENDING_SALTS_SYNC = False
PENDING_RECENTS_SYNC = False
PENDING_OFFLINE_QUEUES_SYNC = False

def load_maintenance():
    if os.path.exists(LOCAL_MAINTENANCE_PATH):
        with open(LOCAL_MAINTENANCE_PATH, "r", encoding="utf-8") as f:
            return json.load(f).get("maintenance", False)
    return False

@app.get("/maintenance")
async def get_maintenance():
    return {"maintenance": load_maintenance()}

@app.get("/")
@app.head("/")
async def health_check():
    return Response(content="OK", media_type="text/plain")

def hash_password(plain_text):
    salt = bcrypt.gensalt()
    return bcrypt.hashpw(plain_text.encode('utf-8'), salt).decode('utf-8')

def verify_password(plain_text, hashed_text):
    try:
        return bcrypt.checkpw(plain_text.encode('utf-8'), hashed_text.encode('utf-8'))
    except Exception:
        return False

file_lock = asyncio.Lock()

def conversation_key(user_a, user_b):
    return "|".join(sorted([user_a, user_b]))

def fetch_hf_file(filename):
    if not hf_hub_url or not HF_TOKEN or not DATASET_REPO_ID:
        raise RuntimeError("Hugging Face storage is not configured")
    url = hf_hub_url(repo_id=DATASET_REPO_ID, filename=filename, repo_type="dataset")
    req = urllib.request.Request(url, headers={"Authorization": f"Bearer {HF_TOKEN}"})
    with urllib.request.urlopen(req) as response:
        return response.read().decode("utf-8")

def load_databases():
    global USERS_DB, CONVERSATIONS_DB, SALTS_DB, RECENTS_DB, OFFLINE_QUEUES_DB, NEXT_USER_ID
    
    # Load offline queues safely
    if os.path.exists(LOCAL_OFFLINE_QUEUES_PATH):
        try:
            with open(LOCAL_OFFLINE_QUEUES_PATH, "r", encoding="utf-8") as f:
                OFFLINE_QUEUES_DB = json.load(f)
        except Exception:
            OFFLINE_QUEUES_DB = {}
    elif HF_TOKEN and DATASET_REPO_ID and hf_hub_download:
        try:
            downloaded_path = hf_hub_download(repo_id=DATASET_REPO_ID, filename=OFFLINE_QUEUES_FILENAME, repo_type="dataset", token=HF_TOKEN)
            with open(downloaded_path, "r", encoding="utf-8") as f:
                OFFLINE_QUEUES_DB = json.load(f)
            with open(LOCAL_OFFLINE_QUEUES_PATH, "w", encoding="utf-8") as f:
                json.dump(OFFLINE_QUEUES_DB, f, indent=4)
        except Exception:
            OFFLINE_QUEUES_DB = {}
    else:
        OFFLINE_QUEUES_DB = {}

    if not HF_TOKEN or not DATASET_REPO_ID:
        print("Running in local-only mode.")
        if os.path.exists(LOCAL_USERS_PATH):
            with open(LOCAL_USERS_PATH, "r", encoding="utf-8") as f:
                raw = json.load(f)
                USERS_DB = raw.get("users", raw)
        else:
            USERS_DB = {}
        if os.path.exists(LOCAL_CONVERSATIONS_PATH):
            with open(LOCAL_CONVERSATIONS_PATH, "r", encoding="utf-8") as f:
                CONVERSATIONS_DB = json.load(f)
        else:
            CONVERSATIONS_DB = {}
        if os.path.exists(LOCAL_SALTS_PATH):
            with open(LOCAL_SALTS_PATH, "r", encoding="utf-8") as f:
                SALTS_DB = json.load(f)
        else:
            SALTS_DB = {}
        if os.path.exists(LOCAL_RECENTS_PATH):
            with open(LOCAL_RECENTS_PATH, "r", encoding="utf-8") as f:
                RECENTS_DB = json.load(f)
        else:
            RECENTS_DB = {}
        existing_ids = [int(u["user_id"]) for u in USERS_DB.values() if str(u.get("user_id", "")).isdigit()]
        NEXT_USER_ID = max(existing_ids, default=0) + 1
        return

    try:
        raw_users = fetch_hf_file(USERS_FILENAME)
        parsed = json.loads(raw_users)
        USERS_DB = parsed.get("users", parsed)
        with open(LOCAL_USERS_PATH, "w", encoding="utf-8") as f:
            json.dump(USERS_DB, f, indent=4)
        print(f"Loaded {len(USERS_DB)} user(s) from HF.")
    except Exception as e:
        print(f"Failed to load users from HF: {e}")
        USERS_DB = {}

    try:
        raw_convos = fetch_hf_file(CONVERSATIONS_FILENAME)
        CONVERSATIONS_DB = json.loads(raw_convos)
        with open(LOCAL_CONVERSATIONS_PATH, "w", encoding="utf-8") as f:
            json.dump(CONVERSATIONS_DB, f, indent=4)
        print(f"Loaded {len(CONVERSATIONS_DB)} conversation(s) from HF.")
    except Exception as e:
        print(f"No conversations file found or failed to load: {e}")
        CONVERSATIONS_DB = {}

    try:
        raw_salts = fetch_hf_file(SALTS_FILENAME)
        SALTS_DB = json.loads(raw_salts)
        with open(LOCAL_SALTS_PATH, "w", encoding="utf-8") as f:
            json.dump(SALTS_DB, f, indent=4)
        print(f"Loaded {len(SALTS_DB)} salt(s) from HF.")
    except Exception as e:
        print(f"No salts file found or failed to load: {e}")
        SALTS_DB = {}

    try:
        raw_recents = fetch_hf_file(RECENTS_FILENAME)
        RECENTS_DB = json.loads(raw_recents)
        with open(LOCAL_RECENTS_PATH, "w", encoding="utf-8") as f:
            json.dump(RECENTS_DB, f, indent=4)
        print(f"loaded {len(RECENTS_DB)} recent(s) from HF.")
    except Exception as e:
        print(f"No recents file or faile to load: {e}")
        RECENTS_DB = {}

    existing_ids = [int(u["user_id"]) for u in USERS_DB.values() if str(u.get("user_id", "")).isdigit()]
    NEXT_USER_ID = max(existing_ids, default=0) + 1

load_databases()

MAX_RECENT_CONVERSATIONS_PER_USER = 100

def update_recent_chat_entry(user, partner, message_text, timestamp, display_name):
    if not user or not partner:
        return
    if user not in RECENTS_DB:
        RECENTS_DB[user] = []

    entry = {
        "partner": partner,
        "last_message": message_text,
        "timestamp": timestamp,
        "display_name": display_name or partner,
    }
    RECENTS_DB[user] = [item for item in RECENTS_DB[user] if item.get("partner") != partner]
    RECENTS_DB[user].append(entry)
    RECENTS_DB[user].sort(key=lambda item: item.get("timestamp", 0), reverse=True)
    RECENTS_DB[user] = RECENTS_DB[user][:MAX_RECENT_CONVERSATIONS_PER_USER]

def save_users_local():
    global PENDING_USERS_SYNC
    with open(LOCAL_USERS_PATH, "w", encoding="utf-8") as f:
        json.dump(USERS_DB, f, indent=4)
    PENDING_USERS_SYNC = True

def save_conversations_local():
    global PENDING_CONVERSATIONS_SYNC
    with open(LOCAL_CONVERSATIONS_PATH, "w", encoding="utf-8") as f:
        json.dump(CONVERSATIONS_DB, f, indent=4)
    PENDING_CONVERSATIONS_SYNC = True

def save_salts_local():
    global PENDING_SALTS_SYNC
    with open(LOCAL_SALTS_PATH, "w", encoding="utf-8") as f:
        json.dump(SALTS_DB, f, indent=4)
    PENDING_SALTS_SYNC = True

def save_recents_local():
    global PENDING_RECENTS_SYNC
    with open(LOCAL_RECENTS_PATH, "w", encoding="utf-8") as f:
        json.dump(RECENTS_DB, f, indent=4)
    PENDING_RECENTS_SYNC = True

def save_offline_queues_local():
    global PENDING_OFFLINE_QUEUES_SYNC
    with open(LOCAL_OFFLINE_QUEUES_PATH, "w", encoding="utf-8") as f:
        json.dump(OFFLINE_QUEUES_DB, f, indent=4)
    PENDING_OFFLINE_QUEUES_SYNC = True

async def background_cloud_sync():
    global PENDING_USERS_SYNC, PENDING_CONVERSATIONS_SYNC, PENDING_SALTS_SYNC, PENDING_RECENTS_SYNC, PENDING_OFFLINE_QUEUES_SYNC
    while True:
        await asyncio.sleep(30)
        if hf_api and DATASET_REPO_ID:
            try:
                if PENDING_USERS_SYNC:
                    await asyncio.to_thread(hf_api.upload_file, path_or_fileobj=LOCAL_USERS_PATH, path_in_repo=USERS_FILENAME, repo_id=DATASET_REPO_ID, repo_type="dataset")
                    PENDING_USERS_SYNC = False
                
                if PENDING_CONVERSATIONS_SYNC:
                    await asyncio.to_thread(hf_api.upload_file, path_or_fileobj=LOCAL_CONVERSATIONS_PATH, path_in_repo=CONVERSATIONS_FILENAME, repo_id=DATASET_REPO_ID, repo_type="dataset")
                    PENDING_CONVERSATIONS_SYNC = False
                    
                if PENDING_SALTS_SYNC:
                    await asyncio.to_thread(hf_api.upload_file, path_or_fileobj=LOCAL_SALTS_PATH, path_in_repo=SALTS_FILENAME, repo_id=DATASET_REPO_ID, repo_type="dataset")
                    PENDING_SALTS_SYNC = False
                    
                if PENDING_RECENTS_SYNC:
                    await asyncio.to_thread(hf_api.upload_file, path_or_fileobj=LOCAL_RECENTS_PATH, path_in_repo=RECENTS_FILENAME, repo_id=DATASET_REPO_ID, repo_type="dataset")
                    PENDING_RECENTS_SYNC = False
                    
                if PENDING_OFFLINE_QUEUES_SYNC:
                    await asyncio.to_thread(hf_api.upload_file, path_or_fileobj=LOCAL_OFFLINE_QUEUES_PATH, path_in_repo=OFFLINE_QUEUES_FILENAME, repo_id=DATASET_REPO_ID, repo_type="dataset")
                    PENDING_OFFLINE_QUEUES_SYNC = False
                    
            except Exception as e:
                print(f"[-] Background cloud sync partial failure: {e}")

@app.on_event("startup")
async def startup_event():
    asyncio.create_task(background_cloud_sync())

ONLINE_USERS = {}
ACTIVE_TOKENS = {}

@app.websocket("/ws")
async def chat_handler(websocket: WebSocket):
    global NEXT_USER_ID
    await websocket.accept()
    username = None
    is_authenticated = False
    try:
        while True:
            raw_data = await websocket.receive_text()
            try:
                auth_data = json.loads(raw_data)
            except json.JSONDecodeError:
                await websocket.send_text("[-] ERROR: Invalid JSON format.")
                continue

            action = auth_data.get("action")
            username_input_raw = auth_data.get("user").lower() if auth_data.get("user") else ""
            username_input_cleaned = username_input_raw.lower()
            password = auth_data.get("pass")
            display_name = auth_data.get("display", username_input_cleaned)

            if action == "signup":
                if not username_input_cleaned or not password:
                    await websocket.send_text("[-] ERROR: Username or password cannot be blank.")
                    continue
                if username_input_cleaned in USERS_DB:
                    await websocket.send_text("[-] FAIL: Username already taken.")
                    continue
                if len(password) >= 8:
                    banned_keywords = ["admin", "websitecreator", "thewebsitecreator", "webs1tecreator", "thewebs1tecreator"]
                    banned_chars = ["{", "}", "(", ")", "[", "]", "@", "#", ";", ":", "<", ">", "&", '"', "'", "\\"]
                    
                    cleaned_display_name = display_name.lower().replace(" ", "")
                    cleaned_username = username_input_cleaned.lower().replace(" ", "")
                    
                    if len(display_name) > 30:
                        await websocket.send_text("[-] ERROR: Display name is too long.")
                        continue
                    elif len(username_input_cleaned) > 30:
                        await websocket.send_text("[-] ERROR: username is too long.")
                        continue
                    elif any(char in display_name for char in banned_chars):
                        await websocket.send_text(r"""[-] ERROR: Display name contains banned character/s. ("{", "}", "(", ")", "[", "]", "@", "#", ";", ":", "<", ">", "&", '"', "'", "\\")""")
                        continue
                    elif any(char in username_input_cleaned for char in banned_chars):
                        await websocket.send_text(r"""[-] ERROR: Username contains banned character/s. ("{", "}", "(", ")", "[", "]", "@", "#", ";", ":", "<", ">", "&", '"', "'", "\\")""")
                        continue
                    elif cleaned_display_name in ["zach(thewebsitecreator)", "zachthewebsitecreator"] or any(item in cleaned_username for item in banned_keywords) or any(item in cleaned_display_name for item in banned_keywords):
                        await websocket.send_text("[-] ERROR: You cannot impersonate me.")
                        continue

                    hashed = await asyncio.to_thread(hash_password, password)
                    async with file_lock:
                        assigned_id = str(NEXT_USER_ID)
                        NEXT_USER_ID += 1
                        USERS_DB[username_input_cleaned] = {
                            "user_id": assigned_id,
                            "password": hashed,
                            "display_name": display_name,
                            "friends": [],
                            "pending_friend_requests": [],
                            "sent_friend_requests": [],
                            "session_token": ""
                        }
                    async with file_lock:
                        SALTS_DB[USERS_DB[username_input_cleaned]["user_id"]] = secrets.token_hex(32)
                    
                    await asyncio.to_thread(save_users_local)
                    await asyncio.to_thread(save_salts_local)
                    await websocket.send_text("[+] SUCCESS: Account created! Please log in.")
                    continue
                else:
                    await websocket.send_text("[-] ERROR: Your password is too short. Please enter a longer password.")
                    continue

            elif action == "token_login":
                token_input = auth_data.get("token")
                remember_token_input = auth_data.get("remember_token")
                matched_user = None
                used_remember_token = False

                if token_input:
                    matched_user = ACTIVE_TOKENS.get(token_input)

                if not matched_user and remember_token_input:
                    for u_key, u_info in USERS_DB.items():
                        u_id = str(u_info.get("user_id"))
                        u_salt = SALTS_DB.get(u_id)
                        stored_hash = u_info.get("remember_token_hash")
                        if u_salt and stored_hash:
                            check_hash = hashlib.sha256((remember_token_input + u_salt).encode()).hexdigest()
                            if check_hash == stored_hash:
                                matched_user = u_key
                                used_remember_token = True
                                break

                if matched_user:
                    username = matched_user
                    if username in ONLINE_USERS:
                        try:
                            await ONLINE_USERS[username].close()
                        except Exception:
                            pass
                    ONLINE_USERS[username] = websocket
                    is_authenticated = True
                    
                    new_token = str(uuid.uuid4())
                    
                    if token_input in ACTIVE_TOKENS:
                        del ACTIVE_TOKENS[token_input]
                    ACTIVE_TOKENS[new_token] = username
                    
                    async with file_lock:
                        USERS_DB[username]["session_token"] = new_token

                    new_remember_token = None
                    if used_remember_token:
                        user_id = str(USERS_DB[username]["user_id"])
                        user_salt = SALTS_DB.get(user_id)
                        if user_salt:
                            new_remember_token = secrets.token_hex(32)
                            new_combined = new_remember_token + user_salt
                            new_salted_hash = hashlib.sha256(new_combined.encode()).hexdigest()
                            async with file_lock:
                                USERS_DB[username]["remember_token_hash"] = new_salted_hash

                    await asyncio.to_thread(save_users_local)

                    token_login_payload = {
                        "action": "login_success",
                        "token": new_token,
                        "user": username,
                        "display_name": USERS_DB[username]["display_name"],
                        "pending_friend_requests": USERS_DB[username]["pending_friend_requests"],
                        "friends": USERS_DB[username]["friends"],
                        "sent_friend_requests": USERS_DB[username].get("sent_friend_requests", []),
                        "recent_chats": RECENTS_DB.get(username, [])
                    }
                    if new_remember_token:
                        token_login_payload["remember_token"] = new_remember_token

                    await websocket.send_text(json.dumps(token_login_payload))
                    
                    if username in OFFLINE_QUEUES_DB:
                        for missed_msgs in OFFLINE_QUEUES_DB[username]:
                            await websocket.send_text(missed_msgs)
                        del OFFLINE_QUEUES_DB[username]
                        await asyncio.to_thread(save_offline_queues_local)
                        
                    print(f"{username} re-authenticated via token.")
                else:
                    await websocket.send_text("[-] FAIL: Invalid or expired session token.")
                continue

            elif action == "login":
                if not username_input_cleaned or not password:
                    await websocket.send_text("[-] FAIL: Username and password required.")
                    continue
                is_valid = False
                if username_input_cleaned in USERS_DB:
                    is_valid = await asyncio.to_thread(verify_password, password, USERS_DB[username_input_cleaned]["password"])
                
                if is_valid:
                    username = username_input_cleaned
                    if username in ONLINE_USERS:
                        try:
                            await ONLINE_USERS[username].close()
                        except Exception:
                            pass
                    ONLINE_USERS[username] = websocket
                    is_authenticated = True
                    
                    token = str(uuid.uuid4())
                    ACTIVE_TOKENS[token] = username
                    
                    user_id = str(USERS_DB[username]["user_id"])
                    remember_me = auth_data.get("remember_me", False)
                    
                    remember_token = None
                    if remember_me:
                        user_salt = SALTS_DB.get(user_id)
                        if user_salt:
                            remember_token = secrets.token_hex(32)
                            combined_string = remember_token + user_salt
                            salted_hash = hashlib.sha256(combined_string.encode()).hexdigest()
                            
                            async with file_lock:
                                USERS_DB[username]["remember_token_hash"] = salted_hash
                            await asyncio.to_thread(save_users_local)
                    
                    async with file_lock:
                        USERS_DB[username]["session_token"] = token
                    await asyncio.to_thread(save_users_local)
                    
                    response_payload = {
                        "action": "login_success", 
                        "token": token, 
                        "user": username,
                        "user_id": user_id,
                        "display_name": USERS_DB[username]["display_name"],
                        "pending_friend_requests": USERS_DB[username]["pending_friend_requests"],
                        "friends": USERS_DB[username]["friends"],
                        "sent_friend_requests": USERS_DB[username].get("sent_friend_requests", []),
                        "recent_chats": RECENTS_DB.get(username, [])
                    }
                    if remember_token:
                        response_payload["remember_token"] = remember_token
                    await websocket.send_text(json.dumps(response_payload))
                    
                    if username in OFFLINE_QUEUES_DB:
                        for missed_msgs in OFFLINE_QUEUES_DB[username]:
                            await websocket.send_text(missed_msgs)
                        del OFFLINE_QUEUES_DB[username]
                        await asyncio.to_thread(save_offline_queues_local)
                        
                    print(f"{username} logged in.")
                    continue
                else:
                    await websocket.send_text("[-] FAIL: Invalid username or password.")
                    continue

            if not is_authenticated:
                await websocket.send_text("[-] ERROR: You must log in or sign up first.")
                continue

            if action == "logout":
                async with file_lock:
                    old_token = USERS_DB[username].get("session_token")
                    if old_token in ACTIVE_TOKENS:
                        del ACTIVE_TOKENS[old_token]
                    USERS_DB[username]["session_token"] = ""
                    USERS_DB[username]["remember_token_hash"] = ""
                ONLINE_USERS.pop(username, None)
                await asyncio.to_thread(save_users_local)
                await websocket.send_text("[+] SUCCESS: Logged out.")
                await websocket.close()
                break

            # Master list for all the functions
            dispatch_table = {
                # I won't be adding comments for these because they're just self-explanatory
                "search_users": search_users,
                "send_friend_request": send_friend_request,
                "accept_friend_request": accept_friend_request,
                "decline_friend_request": decline_friend_request,
                "unfriend": unfriend,
                "send_chat_message": send_chat_message,
                "request_chat_history": request_chat_history,
                "mark_seen": mark_seen,
                "ping": ping
            }

            if action in dispatch_table:
                dispatch_table[action]()
            else:
                await websocket.send_text("[-] ERROR: Invalid action.")
                continue

    except WebSocketDisconnect:
        pass
    except Exception as e:
        print(f"Unhandled error in chat_handler for {username}: {e}")
    finally:
        if username and username in ONLINE_USERS and ONLINE_USERS[username] == websocket:
            del ONLINE_USERS[username]
            print(f"{username} disconnected.")

if __name__ == "__main__":
    import uvicorn
    print("[+] Starting the messenger backend on port 7860...")
    uvicorn.run(app, host="0.0.0.0", port=7860)

the_unreachable_void("Zach", "pass123", "pass123", "Parabola", "No", "NoTokenForYouLol", "pass", "Yes", 1, "Goodbye")