main.py

import os
import json
import secrets
import hashlib
import threading
import requests
from datetime import datetime
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse
from fastapi.exceptions import RequestValidationError
from pydantic import BaseModel, Field
from transformers import pipeline
from huggingface_hub import HfApi, hf_hub_download, create_repo, attach_huggingface_oauth, parse_huggingface_oauth

app = FastAPI()

# Enable official Hugging Face OAuth routes automatically
attach_huggingface_oauth(app)

# Allowed Hugging Face Accounts
ALLOWED_USERS = {
    "MultivexAI", "Datdanboi25", "LH-Tech-AI", "AxionLab-official",
    "Mmorgan-ML", "QyrouNnet-AI", "User01110", "CompactAI",
    "Crownelius", "Enderchef", "armand0e", "Harley-ml"
}

# In-memory cache for profiles to avoid rate limiting
CONTRIBUTORS_CACHE = []

# Fetch permitted user profiles from the HF API in a separate background thread
def fetch_contributor_profiles():
    global CONTRIBUTORS_CACHE
    temp_list = []
    
    with requests.Session() as session:
        for username in ALLOWED_USERS:
            try:
                # Query the public HF User API
                url = f"https://huggingface.co/api/users/{username}/overview"
                response = session.get(url, timeout=5)
                
                if response.ok:
                    data = response.json()
                    fullname = data.get("fullname") or data.get("name") or username
                    avatar_url = data.get("avatarUrl")
                    if not avatar_url:
                        avatar_url = "https://huggingface.co/avatars/default.svg"
                else:
                    fullname = username
                    avatar_url = "https://huggingface.co/avatars/default.svg"
            except Exception as e:
                print(f"[WARNING] Could not retrieve profile for {username}: {str(e)}")
                fullname = username
                avatar_url = "https://huggingface.co/avatars/default.svg"
                
            temp_list.append({
                "username": username,
                "fullname": fullname,
                "avatarUrl": avatar_url
            })
            
    # Atomically update global cache
    CONTRIBUTORS_CACHE = temp_list
    print(f"[INFO] Successfully loaded metadata for {len(CONTRIBUTORS_CACHE)} contributors.")

# Spawn background worker to prevent blocking startup
threading.Thread(target=fetch_contributor_profiles, daemon=True).start()

# Secure connection setup
HF_TOKEN = os.environ.get("HF_TOKEN")
SPACE_ID = os.environ.get("SPACE_ID")

if not HF_TOKEN:
    print("[WARNING] HF_TOKEN is missing. Data persistence disabled.")
    DATASET_REPO_ID = None
else:
    namespace = SPACE_ID.split("/")[0] if SPACE_ID else "local"
    DATASET_REPO_ID = f"{namespace}/atomixlabs-comments-database"
    try:
        api = HfApi(token=HF_TOKEN)
        create_repo(
            repo_id=DATASET_REPO_ID,
            repo_type="dataset",
            private=True,
            exist_ok=True
        )
        print(f"[INFO] Securely connected to private dataset: {DATASET_REPO_ID}")
    except Exception as e:
        print(f"[ERROR] Failed to verify/create private dataset repo: {str(e)}")

db_lock = threading.Lock()
LOCAL_DB_CACHE = "comments_db.json"

# Content Safety Classifier
print("[INFO] Initializing content safety classifier...")
safety_classifier = pipeline("text-classification", model="martin-ha/toxic-comment-model")
print("[INFO] Content safety classifier ready.")

if not os.path.exists(LOCAL_DB_CACHE):
    with open(LOCAL_DB_CACHE, "w") as f:
        json.dump({"welcome": []}, f)

def pull_database():
    if not HF_TOKEN or not DATASET_REPO_ID:
        return load_local_db()
    try:
        filepath = hf_hub_download(
            repo_id=DATASET_REPO_ID,
            filename="comments_db.json",
            repo_type="dataset",
            token=HF_TOKEN
        )
        with open(filepath, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception as e:
        print(f"[INFO] Failed to fetch remote database. Falling back to local. Error: {str(e)}")
        return load_local_db()

def load_local_db():
    if os.path.exists(LOCAL_DB_CACHE):
        with open(LOCAL_DB_CACHE, "r", encoding="utf-8") as f:
            return json.load(f)
    return {"welcome": []}

def push_database(data):
    with open(LOCAL_DB_CACHE, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=4)
    if not HF_TOKEN or not DATASET_REPO_ID:
        return
    try:
        api = HfApi(token=HF_TOKEN)
        api.upload_file(
            path_or_fileobj=LOCAL_DB_CACHE,
            path_in_repo="comments_db.json",
            repo_id=DATASET_REPO_ID,
            repo_type="dataset"
        )
    except Exception as e:
        print(f"[ERROR] Failed to sync data to dataset repo: {str(e)}")

class CommentPayload(BaseModel):
    body: str = Field(..., min_length=1, max_length=1000)

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    return HTMLResponse(status_code=500, content="Internal system integrity error.")

@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    return HTMLResponse(status_code=400, content="Invalid input validation constraints breached.")

# Fetch Login State securely using Hugging Face's OAuth SDK
@app.get("/api/me")
def get_me(request: Request):
    oauth_info = parse_huggingface_oauth(request)
    if oauth_info:
        username = oauth_info.user_info.preferred_username
        is_authorized = any(u.lower() == username.lower() for u in ALLOWED_USERS)
        return {"logged_in": True, "username": username, "authorized": is_authorized}
    return {"logged_in": False}

@app.get("/api/contributors")
def get_contributors():
    # Fallback structure if thread hasn't finished loading yet
    if not CONTRIBUTORS_CACHE:
        return [
            {
                "username": u,
                "fullname": u,
                "avatarUrl": "https://huggingface.co/avatars/default.svg"
            }
            for u in ALLOWED_USERS
        ]
    return CONTRIBUTORS_CACHE

@app.get("/api/comments/{post_id}")
def get_comments(post_id: str):
    db = pull_database()
    comments = db.get(post_id, [])
    
    modified = False
    for c in comments:
        if "id" not in c:
            raw_hash = f"{c.get('username')}_{c.get('date')}_{c.get('body')}"
            c["id"] = hashlib.md5(raw_hash.encode('utf-8')).hexdigest()
            c["edited"] = c.get("edited", False)
            c["history"] = c.get("history", [])
            modified = True
            
    if modified:
        with db_lock:
            current_db = pull_database()
            current_db[post_id] = comments
            push_database(current_db)
            
    return comments

# Post Comment with Hugging Face user validation
@app.post("/api/comments/{post_id}")
def add_comment(post_id: str, payload: CommentPayload, request: Request):
    oauth_info = parse_huggingface_oauth(request)
    if not oauth_info:
        raise HTTPException(status_code=401, detail="Session expired or unauthenticated.")
        
    username = oauth_info.user_info.preferred_username
    is_authorized = any(u.lower() == username.lower() for u in ALLOWED_USERS)
    if not is_authorized:
        raise HTTPException(status_code=403, detail="User account signature rejected.")

    comment_text = payload.body.strip()

    # Content Safety Filter
    prediction = safety_classifier(comment_text)[0]
    if prediction['label'] == 'toxic' and prediction['score'] > 0.75:
         raise HTTPException(
             status_code=400, 
             detail="Content safety screening failure: Toxic pattern detected."
         )

    with db_lock:
        db = pull_database()
        if post_id not in db:
            db[post_id] = []
        
        new_comment = {
            "id": secrets.token_hex(8),
            "username": username,
            "date": datetime.utcnow().strftime("%b %d, %Y, %I:%M %p UTC"),
            "body": comment_text,
            "edited": False,
            "history": []
        }
        db[post_id].append(new_comment)
        push_database(db)
    
    return {"status": "success", "comment": new_comment}

# Edit Comment Endpoint (Verifies ownership and safety filters)
@app.patch("/api/comments/{post_id}/{comment_id}")
def edit_comment(post_id: str, comment_id: str, payload: CommentPayload, request: Request):
    oauth_info = parse_huggingface_oauth(request)
    if not oauth_info:
        raise HTTPException(status_code=401, detail="Session expired or unauthenticated.")
        
    username = oauth_info.user_info.preferred_username
    is_authorized = any(u.lower() == username.lower() for u in ALLOWED_USERS)
    if not is_authorized:
        raise HTTPException(status_code=403, detail="User account signature rejected.")

    comment_text = payload.body.strip()

    # Screen edits for toxicity
    prediction = safety_classifier(comment_text)[0]
    if prediction['label'] == 'toxic' and prediction['score'] > 0.75:
         raise HTTPException(
             status_code=400, 
             detail="Content safety screening failure: Toxic pattern detected in modification."
         )

    with db_lock:
        db = pull_database()
        if post_id not in db:
            raise HTTPException(status_code=404, detail="Post not found.")
            
        comments = db[post_id]
        comment = next((c for c in comments if c.get("id") == comment_id), None)
        
        if not comment:
            raise HTTPException(status_code=404, detail="Comment not found.")
            
        if comment["username"].lower() != username.lower():
            raise HTTPException(status_code=403, detail="Permission denied: Author signature mismatch.")

        # Preserve edit history
        if "history" not in comment:
            comment["history"] = []
            
        comment["history"].append({
            "date": comment.get("last_modified_date", comment["date"]),
            "body": comment["body"]
        })
        
        comment["body"] = comment_text
        comment["edited"] = True
        comment["last_modified_date"] = datetime.utcnow().strftime("%b %d, %Y, %I:%M %p UTC")
        
        push_database(db)
        
    return {"status": "success", "comment": comment}

# Delete Comment Endpoint (Verifies ownership)
@app.delete("/api/comments/{post_id}/{comment_id}")
def delete_comment(post_id: str, comment_id: str, request: Request):
    oauth_info = parse_huggingface_oauth(request)
    if not oauth_info:
        raise HTTPException(status_code=401, detail="Session expired or unauthenticated.")
        
    username = oauth_info.user_info.preferred_username
    is_authorized = any(u.lower() == username.lower() for u in ALLOWED_USERS)
    if not is_authorized:
        raise HTTPException(status_code=403, detail="User account signature rejected.")

    with db_lock:
        db = pull_database()
        if post_id not in db:
            raise HTTPException(status_code=404, detail="Post not found.")
            
        comments = db[post_id]
        comment = next((c for c in comments if c.get("id") == comment_id), None)
        
        if not comment:
            raise HTTPException(status_code=404, detail="Comment not found.")
            
        if comment["username"].lower() != username.lower():
            raise HTTPException(status_code=403, detail="Permission denied: Author signature mismatch.")
            
        db[post_id] = [c for c in comments if c.get("id") != comment_id]
        push_database(db)
        
    return {"status": "success"}

@app.get("/", response_class=HTMLResponse)
async def read_index():
    with open("index.html", "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())