server.py

import gradio as gr
import pandas as pd
import threading
from datetime import datetime
import os
import json
import sqlite3
import time
from dotenv import load_dotenv

DEMO_MODE = os.getenv("DEMO_MODE", "False").lower() == 'true'
load_dotenv()
try:
    from datasets import load_dataset, Dataset, DatasetDict, Features, Value
    HF_DATASETS_AVAILABLE = True
except ImportError:
    HF_DATASETS_AVAILABLE = False
    Features, Value = None, None

STORAGE_BACKEND_CONFIG = os.getenv("STORAGE_BACKEND", "JSON").upper()
HF_DATASET_REPO = os.getenv("HF_DATASET_REPO")
HF_TOKEN = os.getenv("HF_TOKEN")
DB_FILE_JSON = "social_data.json"
DB_FILE_SQLITE = "social_data.db"
db_lock = threading.Lock()
HF_BACKUP_THRESHOLD = int(os.getenv("HF_BACKUP_THRESHOLD", 10))
dirty_operations_count = 0

def force_persist_data():
    global dirty_operations_count
    with db_lock:
        storage_backend = STORAGE_BACKEND_CONFIG
        if storage_backend == "RAM":
            return True, "RAM backend. No persistence."
        elif storage_backend == "SQLITE":
            with sqlite3.connect(DB_FILE_SQLITE) as conn:
                users_df = pd.DataFrame(list(users_db.items()), columns=['username', 'password'])
                users_df.to_sql('users', conn, if_exists='replace', index=False)
                posts_df.to_sql('posts', conn, if_exists='replace', index=False)
                comments_df.to_sql('comments', conn, if_exists='replace', index=False)
            return True, "Successfully saved to SQLite."
        elif storage_backend == "JSON":
            with open(DB_FILE_JSON, "w") as f:
                json.dump({"users": users_db, "posts": posts_df.to_dict('records'), "comments": comments_df.to_dict('records')}, f, indent=2)
            return True, "Successfully saved to JSON file."
        elif storage_backend == "HF_DATASET":
            if not all([HF_DATASETS_AVAILABLE, HF_TOKEN, HF_DATASET_REPO]):
                return False, "HF_DATASET backend is not configured correctly."
            try:
                print("Pushing data to Hugging Face Hub...")
                dataset_dict = DatasetDict({
                    'users': Dataset.from_pandas(pd.DataFrame(list(users_db.items()), columns=['username', 'password'])),
                    'posts': Dataset.from_pandas(posts_df),
                    'comments': Dataset.from_pandas(comments_df)
                })
                dataset_dict.push_to_hub(HF_DATASET_REPO, token=HF_TOKEN, private=True)
                dirty_operations_count = 0
                return True, f"Successfully pushed data to {HF_DATASET_REPO}."
            except Exception as e:
                return False, f"Error pushing to Hugging Face Hub: {e}"
    return False, "Unknown backend."

def handle_persistence_after_change():
    global dirty_operations_count
    storage_backend = STORAGE_BACKEND_CONFIG
    if storage_backend in ["JSON", "SQLITE"]:
        force_persist_data()
    elif storage_backend == "HF_DATASET":
        with db_lock:
            dirty_operations_count += 1
            if dirty_operations_count >= HF_BACKUP_THRESHOLD:
                force_persist_data()

def load_data():
    global STORAGE_BACKEND_CONFIG
    storage_backend = STORAGE_BACKEND_CONFIG
    with db_lock:
        users = {"admin": "password"}
        posts = pd.DataFrame(columns=["post_id", "username", "content", "timestamp"])
        comments = pd.DataFrame(columns=["comment_id", "post_id", "username", "content", "timestamp", "reply_to_comment_id"])

        if storage_backend == "SQLITE":
            try:
                with sqlite3.connect(DB_FILE_SQLITE) as conn:
                    cursor = conn.cursor()
                    cursor.execute("CREATE TABLE IF NOT EXISTS users (username TEXT PRIMARY KEY, password TEXT NOT NULL)")
                    cursor.execute("CREATE TABLE IF NOT EXISTS posts (post_id INTEGER PRIMARY KEY, username TEXT, content TEXT, timestamp TEXT)")
                    cursor.execute("CREATE TABLE IF NOT EXISTS comments (comment_id INTEGER PRIMARY KEY, post_id INTEGER, username TEXT, content TEXT, timestamp TEXT, reply_to_comment_id INTEGER)")
                    cursor.execute("INSERT OR IGNORE INTO users (username, password) VALUES (?, ?)", ("admin", "password"))
                    conn.commit()
                    users = dict(conn.execute("SELECT username, password FROM users").fetchall())
                    posts = pd.read_sql_query("SELECT * FROM posts", conn)
                    comments = pd.read_sql_query("SELECT * FROM comments", conn)
            except Exception as e:
                print(f"CRITICAL: Failed to use SQLite. Falling back to RAM. Error: {e}")
                STORAGE_BACKEND_CONFIG = "RAM"
        elif storage_backend == "JSON":
             if os.path.exists(DB_FILE_JSON):
                try:
                    with open(DB_FILE_JSON, "r") as f: data = json.load(f)
                    users, posts, comments = data.get("users", users), pd.DataFrame(data.get("posts", [])), pd.DataFrame(data.get("comments", []))
                except (json.JSONDecodeError, KeyError): pass
        elif storage_backend == "HF_DATASET":
            if all([HF_DATASETS_AVAILABLE, HF_TOKEN, HF_DATASET_REPO]):
                try:
                    ds_dict = load_dataset(HF_DATASET_REPO, token=HF_TOKEN, trust_remote_code=True)
                    if ds_dict and all(k in ds_dict for k in ['users', 'posts', 'comments']):
                        users = dict(zip(ds_dict['users']['username'], ds_dict['users']['password']))
                        posts = ds_dict['posts'].to_pandas()
                        comments = ds_dict['comments'].to_pandas()
                        print("Successfully loaded data from HF Dataset.")
                    else:
                        raise ValueError("Dataset dictionary is empty or malformed.")
                except Exception as e:
                    print(f"Could not load from HF Dataset '{HF_DATASET_REPO}'. Attempting to initialize. Error: {e}")
                    try:
                        user_features = Features({'username': Value('string'), 'password': Value('string')})
                        post_features = Features({'post_id': Value('int64'), 'username': Value('string'), 'content': Value('string'), 'timestamp': Value('string')})
                        comment_features = Features({'comment_id': Value('int64'), 'post_id': Value('int64'), 'username': Value('string'), 'content': Value('string'), 'timestamp': Value('string'), 'reply_to_comment_id': Value('int64')})
                        
                        initial_users_df = pd.DataFrame(list(users.items()), columns=['username', 'password'])
                        
                        dataset_dict = DatasetDict({
                            'users': Dataset.from_pandas(initial_users_df, features=user_features),
                            'posts': Dataset.from_pandas(posts, features=post_features),
                            'comments': Dataset.from_pandas(comments, features=comment_features)
                        })
                        dataset_dict.push_to_hub(HF_DATASET_REPO, token=HF_TOKEN, private=True)
                        print(f"Successfully initialized new empty HF Dataset at {HF_DATASET_REPO}.")
                    except Exception as e_push:
                        print(f"CRITICAL: Failed to create new HF Dataset. Falling back to RAM. Push Error: {e_push}")
                        STORAGE_BACKEND_CONFIG = "RAM"
            else:
                 print("HF_DATASET backend not fully configured. Falling back to RAM.")
                 STORAGE_BACKEND_CONFIG = "RAM"

    if "reply_to_comment_id" not in comments.columns:
        comments["reply_to_comment_id"] = None
        
    post_counter = int(posts['post_id'].max()) if not posts.empty else 0
    comment_counter = int(comments['comment_id'].max()) if not comments.empty else 0
    return users, posts, comments, post_counter, comment_counter

users_db, posts_df, comments_df, post_counter, comment_counter = load_data()

def api_register(username, password):
    if not username or not password: return "Failed: Username/password cannot be empty."
    with db_lock:
        if username in users_db: return f"Failed: Username '{username}' already exists."
        users_db[username] = password
        handle_persistence_after_change()
    return f"Success: User '{username}' registered."

def api_login(username, password):
    return f"{username}:{password}" if users_db.get(username) == password else "Failed: Invalid credentials."

def _get_user_from_token(token):
    if not token or ':' not in token: return None
    user, pwd = token.split(':', 1)
    return user if users_db.get(user) == pwd else None

def api_create_post(auth_token, content):
    global posts_df, post_counter
    username = _get_user_from_token(auth_token)
    if not username: return "Failed: Invalid auth token."
    with db_lock:
        post_counter += 1
        new_post = pd.DataFrame([{"post_id": post_counter, "username": username, "content": content, "timestamp": datetime.utcnow().isoformat()}])
        posts_df = pd.concat([posts_df, new_post], ignore_index=True)
        handle_persistence_after_change()
    return f"Success: Post {post_counter} created."

def api_create_comment(auth_token, post_id, content, reply_to_comment_id=None):
    global comments_df, comment_counter
    username = _get_user_from_token(auth_token)
    if not username: return "Failed: Invalid auth token."
    with db_lock:
        if int(post_id) not in posts_df['post_id'].values: return f"Failed: Post {post_id} not found."
        if reply_to_comment_id is not None and int(reply_to_comment_id) not in comments_df['comment_id'].values: return f"Failed: Comment to reply to ({reply_to_comment_id}) not found."
        
        comment_counter += 1
        new_comment = pd.DataFrame([{"comment_id": comment_counter, "post_id": int(post_id), "username": username, "content": content, "timestamp": datetime.utcnow().isoformat(), "reply_to_comment_id": int(reply_to_comment_id) if reply_to_comment_id is not None else None}])
        comments_df = pd.concat([comments_df, new_comment], ignore_index=True)
        handle_persistence_after_change()
    return "Success: Comment created."

def api_get_feed():
    with db_lock:
        posts, comments = posts_df.copy(), comments_df.copy()

    if posts.empty and comments.empty:
        return pd.DataFrame(columns=['type', 'post_id', 'comment_id', 'reply_to_comment_id', 'username', 'timestamp', 'content'])

    posts['type'] = 'post'
    comments['type'] = 'comment'
    
    feed_data = pd.concat([posts, comments], ignore_index=True, sort=False)
    feed_data['timestamp'] = pd.to_datetime(feed_data['timestamp'])
    
    feed_data = feed_data.sort_values(by=['timestamp'], ascending=False)
    
    display_columns = ['type', 'post_id', 'comment_id', 'reply_to_comment_id', 'username', 'timestamp', 'content']
    feed_data = feed_data.reindex(columns=display_columns)
    
    return feed_data.fillna('')

def ui_manual_post(username, password, content):
    auth_token = api_login(username, password)
    if "Failed" in auth_token: return "Login failed.", api_get_feed()
    return api_create_post(auth_token, content), api_get_feed()

def ui_manual_comment(username, password, post_id, reply_id, content):
    auth_token = api_login(username, password)
    if "Failed" in auth_token: return "Login failed.", api_get_feed()
    return api_create_comment(auth_token, post_id, content, reply_id), api_get_feed()

with gr.Blocks(theme=gr.themes.Soft(), title="Social App") as demo:
    gr.Markdown("# Social Media Server for iLearn Agent")
    gr.Markdown(f"This app provides an API for iLearn agents to interact with. **Storage Backend: `{STORAGE_BACKEND_CONFIG}`**")
    
    with gr.Tabs():
        with gr.TabItem("Live Feed"):
            feed_df_display = gr.DataFrame(label="Feed", interactive=False, wrap=True)
            refresh_btn = gr.Button("Refresh Feed")
        
        with gr.TabItem("Manual Actions"):
            manual_action_status = gr.Textbox(label="Action Status", interactive=False)
            with gr.Row():
                with gr.Group():
                    gr.Markdown("### Create Post")
                    post_user = gr.Textbox(label="User", value="admin")
                    post_pass = gr.Textbox(label="Pass", type="password", value="password")
                    post_content = gr.Textbox(label="Content", lines=3)
                    post_button = gr.Button("Submit Post", variant="primary")
                with gr.Group():
                    gr.Markdown("### Create Comment / Reply")
                    comment_user = gr.Textbox(label="User", value="admin")
                    comment_pass = gr.Textbox(label="Pass", type="password", value="password")
                    comment_post_id = gr.Number(label="Target Post ID")
                    comment_reply_id = gr.Number(label="Reply to Comment ID (optional)")
                    comment_content = gr.Textbox(label="Content", lines=2)
                    comment_button = gr.Button("Submit Comment", variant="primary")

    post_button.click(ui_manual_post, [post_user, post_pass, post_content], [manual_action_status, feed_df_display])
    comment_button.click(ui_manual_comment, [comment_user, comment_pass, comment_post_id, comment_reply_id, comment_content], [manual_action_status, feed_df_display])
    refresh_btn.click(api_get_feed, None, feed_df_display)
    
    demo.load(api_get_feed, None, feed_df_display)

    with gr.Column(visible=False):
        gr.Interface(api_register, ["text", "text"], "text", api_name="register")
        gr.Interface(api_login, ["text", "text"], "text", api_name="login")
        gr.Interface(api_create_post, ["text", "text"], "text", api_name="create_post")
        gr.Interface(api_create_comment, ["text", "number", "text", "number"], "text", api_name="create_comment")
        gr.Interface(api_get_feed, None, "dataframe", api_name="get_feed")

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=False)