import os
import uuid
import requests
import io
import logging
from flask import Flask, request, jsonify, send_from_directory
from huggingface_hub import HfApi
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")
HF_REPO_ID = "Asrasahar/Ezmary-uploads"

# 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"Public File Upload API Initialized. Target repo: {HF_REPO_ID}")

# --- توابع کمکی ---
def get_folder_by_mimetype(mimetype):
    if not mimetype: return "others"
    main_type = mimetype.split('/')[0]
    if main_type == "image": return "images"
    elif main_type == "video": return "videos"
    elif main_type == "audio": return "audios"
    elif main_type == "application":
        if 'pdf' in mimetype: return "documents/pdfs"
        elif 'zip' in mimetype or 'rar' in mimetype or '7z' in mimetype: return "archives"
        else: return "documents/applications"
    elif main_type == "text": return "documents/texts"
    else: return "others"

def process_and_upload(file_content: bytes, original_filename: str, mimetype: str) -> str:
    file_extension = os.path.splitext(original_filename)[1] or mimetypes.guess_extension(mimetype) or ""
    filename = f"{uuid.uuid4().hex}{file_extension}"
    target_folder = get_folder_by_mimetype(mimetype)
    path_in_repo = f"{target_folder}/{filename}"
    
    file_stream = io.BytesIO(file_content)
    logger.info(f"Uploading to Hugging Face: {path_in_repo}")
    hf_api.upload_file(path_or_fileobj=file_stream, path_in_repo=path_in_repo, repo_id=HF_REPO_ID, repo_type="dataset")
    direct_url = f"https://huggingface.co/datasets/{HF_REPO_ID}/resolve/main/{path_in_repo}"
    logger.info(f"File uploaded successfully! URL: {direct_url}")
    return direct_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:
        # حالت اول: آپلود از کامپیوتر یا ربات (حالت FormData)
        if 'file' in request.files:
            uploaded_file = request.files['file']
            if uploaded_file.filename == '':
                return jsonify({"error": "No selected file"}), 400
            file_content = uploaded_file.read()
            direct_url = process_and_upload(file_content, uploaded_file.filename, uploaded_file.mimetype)
            return jsonify({"message": "File uploaded successfully!", "hf_url": direct_url}), 200

        # ***** این بخش دوباره اضافه شد *****
        # حالت دوم: آپلود از طریق لینک در وب سایت (حالت JSON)
        elif request.is_json and 'url' in request.json:
            file_url = request.json['url']
            if not file_url:
                return jsonify({"error": "No URL provided"}), 400

            logger.info(f"Downloading file from URL: {file_url}")
            response = requests.get(file_url, stream=True, timeout=30)
            response.raise_for_status()
            
            file_content = response.content
            mimetype = response.headers.get('Content-Type')
            original_filename = file_url.split('/')[-1].split('?')[0]

            direct_url = process_and_upload(file_content, original_filename, mimetype)
            return jsonify({"message": "File uploaded successfully!", "hf_url": direct_url}), 200
        # ***** پایان بخش اضافه شده *****
            
        else:
            return jsonify({"error": "Invalid request. Expected 'file' (form-data) or 'url' (json)."}), 400

    except requests.exceptions.RequestException as e:
        logger.error(f"Failed to download from URL: {e}", exc_info=True)
        return jsonify({"error": f"Failed to download from URL: {str(e)}"}), 400
    except Exception as e:
        logger.error(f"An internal server error occurred: {e}", exc_info=True)
        return jsonify({"error": f"An internal server error occurred: {str(e)}"}), 500

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