import os
import json
import time
import threading
import logging
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import RepositoryNotFoundError, EntryNotFoundError

# --- تنظیمات ---
DATASET_REPO = "ezmarynoori/Karbaran-rayegan-tedad"
DATASET_FILENAME = "voice_conversion_usage_data.json"
USAGE_LIMIT_PER_MODEL = 3 # محدودیت ۳ عدد برای هر مدل در هفته
HF_TOKEN = os.environ.get("HF_TOKEN")

# پوشه بیلد شده توسط Docker (dist)
STATIC_FOLDER = 'dist'

# --- راه‌اندازی برنامه ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
app = Flask(__name__, static_folder=STATIC_FOLDER, static_url_path='')
CORS(app)

# --- متغیرهای ذخیره‌سازی داده ---
usage_data_cache = []
cache_lock = threading.Lock()
data_changed = threading.Event()
api = None

if not HF_TOKEN:
    logging.error("هشدار: توکن HF_TOKEN یافت نشد. دیتابیس کار نخواهد کرد.")
else:
    api = HfApi(token=HF_TOKEN)

# --- لود اولیه داده‌ها ---
def load_initial_data():
    global usage_data_cache
    with cache_lock:
        if not api: return
        try:
            local_path = hf_hub_download(
                repo_id=DATASET_REPO, 
                filename=DATASET_FILENAME, 
                repo_type="dataset",
                token=HF_TOKEN,
                force_download=True
            )
            with open(local_path, 'r', encoding='utf-8') as f:
                content = f.read()
                if content:
                    usage_data_cache = json.loads(content)
        except (RepositoryNotFoundError, EntryNotFoundError):
            usage_data_cache = []
        except Exception as e:
            logging.warning(f"Error loading data: {e}")
            usage_data_cache = []

# --- ذخیره‌سازی در پس‌زمینه ---
def persist_data_to_hub():
    if not api: return
    
    data_snapshot = None
    with cache_lock:
        if data_changed.is_set():
            data_snapshot = list(usage_data_cache)
            data_changed.clear()
    
    if data_snapshot is not None:
        temp_filepath = "/tmp/temp_voice_usage.json"
        try:
            with open(temp_filepath, 'w', encoding='utf-8') as f:
                json.dump(data_snapshot, f, ensure_ascii=False, indent=2)
            
            api.upload_file(
                path_or_fileobj=temp_filepath,
                path_in_repo=DATASET_FILENAME,
                repo_id=DATASET_REPO,
                repo_type="dataset",
                commit_message="Update voice usage"
            )
            logging.info("Data saved to HF Hub.")
        except Exception as e:
            logging.error(f"Failed to save data: {e}")
            data_changed.set()

def background_persister():
    while True:
        time.sleep(300) # هر 5 دقیقه ذخیره کن
        persist_data_to_hub()

def get_user_identifier(data):
    fingerprint = data.get('fingerprint')
    if fingerprint: return str(fingerprint)
    if request.headers.getlist("X-Forwarded-For"):
        return request.headers.getlist("X-Forwarded-For")[0].split(',')[0].strip()
    return request.remote_addr

# --- API Endpoints ---

@app.route('/')
def index():
    return send_from_directory(STATIC_FOLDER, 'index.html')

@app.route('/<path:path>')
def serve_static(path):
    if os.path.exists(os.path.join(STATIC_FOLDER, path)):
        return send_from_directory(STATIC_FOLDER, path)
    return send_from_directory(STATIC_FOLDER, 'index.html')

@app.route('/api/check-credit', methods=['POST'])
def check_credit():
    data = request.get_json()
    if not data: return jsonify({"error": "Invalid request"}), 400
    
    user_id = get_user_identifier(data)
    model_id = data.get('model_id', 'custom')
    
    with cache_lock:
        now = time.time()
        one_week = 7 * 24 * 60 * 60
        user_record = next((u for u in usage_data_cache if u.get('id') == user_id), None)
        
        credits_remaining = USAGE_LIMIT_PER_MODEL
        limit_reached = False
        reset_timestamp = 0
        
        if user_record:
            # بررسی انقضای هفته برای کل کاربر (ریست همه مدل‌ها)
            if user_record.get('week_start', 0) < (now - one_week):
                user_record['usage'] = {} 
                user_record['week_start'] = now
                data_changed.set()
            
            usage_dict = user_record.get('usage', {})
            current_model_usage = usage_dict.get(model_id, 0)
            
            credits_remaining = max(0, USAGE_LIMIT_PER_MODEL - current_model_usage)
            
            if credits_remaining == 0:
                limit_reached = True
                reset_timestamp = user_record.get('week_start', now) + one_week
                
    return jsonify({
        "credits_remaining": credits_remaining,
        "limit_reached": limit_reached,
        "reset_timestamp": reset_timestamp,
        "model_id": model_id
    })

@app.route('/api/use-credit', methods=['POST'])
def use_credit():
    data = request.get_json()
    if not data: return jsonify({"error": "Invalid request"}), 400
    
    user_id = get_user_identifier(data)
    model_id = data.get('model_id', 'custom')
    
    with cache_lock:
        now = time.time()
        one_week = 7 * 24 * 60 * 60
        user_record = next((u for u in usage_data_cache if u.get('id') == user_id), None)
        
        if user_record:
            # ریست هفتگی
            if user_record.get('week_start', 0) < (now - one_week):
                user_record['usage'] = {}
                user_record['week_start'] = now
            
            if 'usage' not in user_record:
                user_record['usage'] = {}
                
            current_model_usage = user_record['usage'].get(model_id, 0)
            
            if current_model_usage >= USAGE_LIMIT_PER_MODEL:
                reset_timestamp = user_record.get('week_start', now) + one_week
                return jsonify({
                    "status": "limit_reached", 
                    "credits_remaining": 0,
                    "reset_timestamp": reset_timestamp
                }), 429
            
            user_record['usage'][model_id] = current_model_usage + 1
        else:
            # کاربر جدید
            user_record = {
                "id": user_id, 
                "week_start": now, 
                "usage": {model_id: 1}
            }
            usage_data_cache.append(user_record)
        
        credits_remaining = USAGE_LIMIT_PER_MODEL - user_record['usage'][model_id]
        data_changed.set()
        
    return jsonify({"status": "success", "credits_remaining": credits_remaining})

# --- Main ---
if __name__ != '__main__':
    load_initial_data()
    t = threading.Thread(target=background_persister, daemon=True)
    t.start()

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)