import os
import uuid
import requests
import io
import logging
import sqlite3
import time
from datetime import timedelta
from flask import Flask, request, jsonify, send_from_directory, Response
from huggingface_hub import HfApi
from apscheduler.schedulers.background import BackgroundScheduler
import mimetypes

# 1. تنظیمات لاگینگ
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger(__name__)

app = Flask(__name__, static_folder='static', static_url_path='/static')

# 2. خواندن متغیرهای محیطی و تعریف ثابت‌ها
HF_TOKEN = os.getenv("HF_TOKEN")
DB_FILE = "/tmp/uploads.db"

space_host = os.getenv('SPACE_HOST')
if space_host:
    UPLOADER_BASE_URL = f"https://{space_host}"
else:
    port = os.getenv('PORT', 7860)
    UPLOADER_BASE_URL = f"http://localhost:{port}"

HF_REPO_IDS = [
    "Hamed744/Uploadfile1",
    "Hamed744/Uploadfile2",
    "Hamed744/Uploadfile3",
    "Hamed744/Uploadfile4",
    "Hamed744/Uploadfile5"
]

# 3. بررسی وجود توکن
if not HF_TOKEN:
    logger.error("FATAL ERROR: HF_TOKEN environment variable not set!")
    exit()

# 4. ایجاد یک نمونه از HfApi
hf_api = HfApi(token=HF_TOKEN)
logger.info(f"File Upload API Initialized. Auto-detected Base URL: {UPLOADER_BASE_URL}")
logger.info(f"Target repos: {', '.join(HF_REPO_IDS)}")
logger.info(f"Using database file at: {DB_FILE}")

# --- بخش مدیریت پایگاه داده و توابع کمکی ---
# (این بخش‌ها بدون تغییر باقی می‌مانند)
def init_db():
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute('CREATE TABLE IF NOT EXISTS uploaded_files (id INTEGER PRIMARY KEY, path_in_repo TEXT NOT NULL, repo_id TEXT NOT NULL, upload_timestamp REAL NOT NULL)')
    cursor.execute('CREATE TABLE IF NOT EXISTS app_state (key TEXT PRIMARY KEY, value INTEGER NOT NULL)')
    cursor.execute("INSERT OR IGNORE INTO app_state (key, value) VALUES ('repo_index', 0)")
    conn.commit()
    conn.close()
    logger.info("Database initialized successfully.")

init_db()

def get_next_repo():
    with sqlite3.connect(DB_FILE) as conn:
        cursor = conn.cursor()
        cursor.execute("SELECT value FROM app_state WHERE key = 'repo_index'")
        current_index = cursor.fetchone()[0]
        repo_id = HF_REPO_IDS[current_index]
        next_index = (current_index + 1) % len(HF_REPO_IDS)
        cursor.execute("UPDATE app_state SET value = ? WHERE key = 'repo_index'", (next_index,))
        conn.commit()
        logger.info(f"Selected repo for current upload: {repo_id} (next index will be {next_index})")
        return repo_id

def add_file_to_db(path_in_repo, repo_id):
    with sqlite3.connect(DB_FILE) as conn:
        conn.execute("INSERT INTO uploaded_files (path_in_repo, repo_id, upload_timestamp) VALUES (?, ?, ?)",
                     (path_in_repo, repo_id, time.time()))
        conn.commit()
    logger.info(f"File metadata added to DB: {path_in_repo} in repo {repo_id}")

def delete_old_files():
    # این تابع کامل است و نیازی به تغییر ندارد
    with app.app_context():
        logger.info("Running scheduled job: Checking for old files to delete...")
        two_weeks_ago = time.time() - timedelta(days=14).total_seconds()
        try:
            with sqlite3.connect(DB_FILE) as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT id, path_in_repo, repo_id FROM uploaded_files WHERE upload_timestamp < ?", (two_weeks_ago,))
                for file_id, path, repo_id in cursor.fetchall():
                    try:
                        logger.warning(f"Deleting old file: {path} from {repo_id}")
                        hf_api.delete_file(path_in_repo=path, repo_id=repo_id, repo_type="dataset")
                        cursor.execute("DELETE FROM uploaded_files WHERE id = ?", (file_id,))
                        conn.commit()
                        logger.info(f"Successfully deleted {path}")
                    except Exception as e:
                        logger.error(f"Failed to delete {path}: {e}")
        except Exception as e:
            logger.error(f"DB error in deletion job: {e}")

def get_folder_by_mimetype(mimetype):
    if not mimetype: return "others"
    main_type = mimetype.split('/')[0]
    types = {"image": "images", "video": "videos", "audio": "audios", "text": "documents/texts"}
    if main_type in types: return types[main_type]
    if main_type == "application":
        if 'pdf' in mimetype: return "documents/pdfs"
        if any(x in mimetype for x in ['zip', 'rar', '7z']): return "archives"
        return "documents/applications"
    return "others"

def process_and_upload(file_content, original_filename, mimetype):
    target_repo_id = get_next_repo()
    ext = os.path.splitext(original_filename)[1] or mimetypes.guess_extension(mimetype) or ""
    filename = f"{uuid.uuid4().hex}{ext}"
    folder = get_folder_by_mimetype(mimetype)
    path_in_repo = f"{folder}/{filename}"
    hf_api.upload_file(
        path_or_fileobj=io.BytesIO(file_content),
        path_in_repo=path_in_repo,
        repo_id=target_repo_id,
        repo_type="dataset"
    )
    add_file_to_db(path_in_repo, target_repo_id)
    # *** مهم: نحوه ساخت URL تغییری نمی‌کند، چون target_repo_id خودش شامل user/repo است ***
    proxy_url = f"{UPLOADER_BASE_URL}/proxy/{target_repo_id}/{path_in_repo}"
    logger.info(f"Full Proxy URL generated: {proxy_url}")
    return proxy_url

# --- بخش API ---
@app.route('/')
def serve_html():
    return send_from_directory(app.static_folder, 'index.html')

@app.route('/upload', methods=['POST'])
def upload_file_endpoint():
    try:
        if 'file' in request.files:
            file = request.files['file']
            if not file.filename: return jsonify({"error": "No selected file"}), 400
            proxy_url = process_and_upload(file.read(), file.filename, file.mimetype)
            return jsonify({"hf_url": proxy_url}), 200
        elif request.is_json and 'url' in request.json:
            url = request.json['url']
            if not url: return jsonify({"error": "No URL provided"}), 400
            resp = requests.get(url, timeout=60)
            resp.raise_for_status()
            filename = url.split('/')[-1].split('?')[0]
            proxy_url = process_and_upload(resp.content, filename, resp.headers.get('Content-Type'))
            return jsonify({"hf_url": proxy_url}), 200
        return jsonify({"error": "Invalid request"}), 400
    except Exception as e:
        logger.error(f"Internal server error: {e}", exc_info=True)
        return jsonify({"error": str(e)}), 500

# *** رفع خطا: تغییر ساختار Route برای پراکسی ***
@app.route('/proxy/<user>/<repo>/<path:file_path>')
def file_proxy(user, repo, file_path):
    repo_id = f"{user}/{repo}" # بازسازی repo_id کامل
    
    if not HF_TOKEN:
        logger.error("FATAL ERROR: HF_TOKEN not configured for proxy!")
        return "Server token not configured", 500

    file_url = f"https://huggingface.co/datasets/{repo_id}/resolve/main/{file_path}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    
    try:
        hf_response = requests.get(file_url, headers=headers, stream=True, timeout=30)
        hf_response.raise_for_status()
        return Response(
            hf_response.iter_content(chunk_size=8192),
            content_type=hf_response.headers.get('Content-Type', 'application/octet-stream')
        )
    except requests.exceptions.HTTPError as e:
         if e.response.status_code == 404:
            logger.error(f"File not found on Hugging Face: {file_url}")
            return "File not found or access denied", 404
         raise e
    except requests.exceptions.RequestException as e:
        logger.error(f"Proxy request error for '{file_url}': {e}")
        return "Error fetching file", 500

# --- راه‌اندازی زمان‌بندی و برنامه ---
scheduler = BackgroundScheduler(daemon=True)
scheduler.add_job(delete_old_files, 'interval', hours=24)
scheduler.start()

if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=False)