import os
import io
import time
import requests
import base64
import threading
from flask import Flask, request, jsonify, render_template_string
from PIL import Image, ImageOps
from huggingface_hub import HfApi, CommitOperationAdd

app = Flask(__name__)

# --- ⚙️ SECRETS (Environment Variables) ---
HF_TOKEN = os.environ.get("HF_TOKEN")
HF_REPO = os.environ.get("HF_REPO") 
WORKER_URL = os.environ.get("WORKER_URL", "https://your-worker.subdomain.workers.dev/") # CF Worker URL
XOR_KEY = os.environ.get("XOR_KEY", "badal_master_key") # Tumhari Secret Chabi 🔑

Image.MAX_IMAGE_PIXELS = None

# --- 📁 LOCAL STORAGE SETUP (Queue ke liye) ---
PENDING_DIR = "pending_uploads"
os.makedirs(PENDING_DIR, exist_ok=True)

# --- 🔄 BACKGROUND BATCH UPLOADER (Har 30 seconds) ---
def background_uploader():
    api = HfApi()
    while True:
        time.sleep(30) # ⏱️ 30 Second Timer (Tune bola tha)
        try:
            files_to_upload = os.listdir(PENDING_DIR)
            if not files_to_upload:
                continue
                
            print(f"🚀 [BATCH UPLOAD] {len(files_to_upload)} files ko ek saath upload kar rahe hain...")
            
            operations = []
            uploaded_files = []
            
            for filename in files_to_upload:
                file_path = os.path.join(PENDING_DIR, filename)
                
                # 🗂️ 10k LIMIT BYPASS: Smart Folder System (Max ~9000 safe capacity)
                # Naam (img_177237_3a7b9c.webp) me se '3a7' nikalega. Isse 4096 alag folders banenge!
                try:
                    sub_folder = filename.split('_')[2][:3] 
                except:
                    sub_folder = "misc"
                    
                operations.append(
                    CommitOperationAdd(
                        # Ab file 'posters/3a7/img_xxx.webp' me jayegi
                        path_in_repo=f"posters/{sub_folder}/{filename}",
                        path_or_fileobj=file_path
                    )
                )
                uploaded_files.append(file_path)
            
            # Single Commit on Hugging Face
            api.create_commit(
                repo_id=HF_REPO,
                repo_type="dataset",
                token=HF_TOKEN,
                operations=operations,
                commit_message=f"Batch upload of {len(operations)} images by OptiPix (Sharded)"
            )
            print("✅ [BATCH UPLOAD] Success! Local files delete kar rahe hain...")
            
            for f in uploaded_files:
                try:
                    os.remove(f)
                except:
                    pass
        except Exception as e:
            print(f"❌ [BATCH UPLOAD] Error: {e}")

# App start hote hi background thread on
threading.Thread(target=background_uploader, daemon=True).start()

# --- 🔐 XOR ENCRYPTION ENGINE ---
def xor_encrypt_decrypt_bytes(data: bytes, key: str) -> bytes:
    """Image ke binary data ko corrupt/decrypt karta hai"""
    key_bytes = key.encode()
    key_len = len(key_bytes)
    return bytes([b ^ key_bytes[i % key_len] for i, b in enumerate(data)])

def encrypt_text(text: str, key: str) -> str:
    """Ab ye sidha target URL ko encrypt karega"""
    xored = "".join(chr(ord(c) ^ ord(key[i % len(key)])) for i, c in enumerate(text))
    return base64.urlsafe_b64encode(xored.encode()).decode().rstrip("=")

# --- 🧠 SUPER-OPTIMIZER LOGIC ---
def dynamic_optimize(image_content, level="extreme"):
    """Image ko optimize karta hai"""
    if level == "none":
        return image_content

    try:
        img = Image.open(io.BytesIO(image_content))
        img = ImageOps.exif_transpose(img)
        if img.mode in ("RGBA", "P"): img = img.convert("RGB")
        
        settings = {
            "extreme": {"width": 550, "quality": 70, "method": 6, "target_kb": 60},
            "high":    {"width": 650, "quality": 80, "method": 5, "target_kb": 90},
            "medium":  {"width": 800, "quality": 85, "method": 4, "target_kb": 150},
            "low":     {"width": 1000, "quality": 95, "method": 3, "target_kb": 300}
        }
        conf = settings.get(level, settings["extreme"])
        
        if img.width > conf["width"]:
            ratio = conf["width"] / float(img.width)
            img = img.resize((conf["width"], int(float(img.height) * ratio)), Image.Resampling.LANCZOS)
            
        output = io.BytesIO()
        curr_q = conf["quality"]
        
        while curr_q > 10:
            output.seek(0)
            output.truncate(0)
            img.save(output, format='WEBP', quality=curr_q, method=conf["method"], optimize=True)
            if (output.tell() / 1024) <= conf["target_kb"]: break
            curr_q -= 5

        return output.getvalue()
    except Exception as e:
        print(f"❌ Optimization Error: {e}")
        return image_content

# --- 🌐 API ROUTE (Handles BOTH URL and DIRECT FILE) ---
@app.route('/upload-poster', methods=['GET', 'POST'])
def handle_upload():
    comp_level = request.args.get('level', request.form.get('level', 'extreme')).lower()
    image_content = None

    if not HF_TOKEN or not HF_REPO: 
        return jsonify({"success": False, "error": "HF Secrets missing!"}), 500

    try:
        if 'file' in request.files and request.files['file'].filename != '':
            file = request.files['file']
            image_content = file.read()
        
        elif request.args.get('url') or request.form.get('url'):
            img_url = request.args.get('url') or request.form.get('url')
            resp = requests.get(img_url, timeout=20)
            if resp.status_code != 200: 
                return jsonify({"success": False, "error": "URL Download failed"}), 400
            image_content = resp.content
            
        if not image_content:
            return jsonify({"success": False, "error": "No file or URL provided"}), 400
        
        # Optimize & Encrypt Data
        optimized_data = dynamic_optimize(image_content, comp_level)
        encrypted_image_data = xor_encrypt_decrypt_bytes(optimized_data, XOR_KEY)
        
        # Generate Unique Filename
        hex_random = os.urandom(3).hex()
        real_filename = f"img_{int(time.time())}_{hex_random}.webp"
        
        # 🗂️ Calculate Sub-folder based on logic
        sub_folder = hex_random[:3]
        
        # 🌐 ASLI JADU: Poora Target URL banao
        full_hf_url = f"https://huggingface.co/datasets/{HF_REPO}/resolve/main/posters/{sub_folder}/{real_filename}"
        
        # 🔐 URL ko encrypt karo taaki worker ko de sake
        encrypted_url_payload = encrypt_text(full_hf_url, XOR_KEY)
        
        # Save to Local Queue
        local_filepath = os.path.join(PENDING_DIR, real_filename)
        with open(local_filepath, "wb") as f:
            f.write(encrypted_image_data)
        
        # Create Final Worker URL
        worker_base = WORKER_URL if WORKER_URL.endswith('/') else WORKER_URL + '/'
        final_cf_url = f"{worker_base}{encrypted_url_payload}"
        
        return jsonify({
            "success": True, 
            "url": final_cf_url,
            "real_hf_filename": real_filename, 
            "level_used": comp_level,
            "original_size_kb": round(len(image_content)/1024, 2),
            "encrypted_size_kb": round(len(encrypted_image_data)/1024, 2)
        })
        
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500

# --- 🖥️ UI ROUTE (TESTING DASHBOARD) ---
HTML_UI = """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>OptiPix Pro | Smart Image Optimizer</title>
    <style>
        :root { --primary: #3b82f6; --bg: #0f172a; --card: #1e293b; --text: #f1f5f9; }
        body { font-family: 'Inter', system-ui, sans-serif; background: var(--bg); color: var(--text); padding: 20px; max-width: 700px; margin: auto; }
        .card { background: var(--card); padding: 30px; border-radius: 16px; box-shadow: 0 10px 25px rgba(0,0,0,0.3); border: 1px solid #334155; }
        h2 { color: #fff; margin-top: 0; font-size: 24px; text-align: center; }
        p.sub { text-align: center; color: #94a3b8; font-size: 14px; margin-bottom: 25px; }
        .input-group { background: #0f172a; padding: 15px; border-radius: 10px; margin-bottom: 20px; border: 1px solid #334155; }
        label { display: block; font-weight: 600; margin-bottom: 8px; font-size: 14px; color: #3b82f6; }
        input[type="text"], input[type="file"], select { width: 100%; padding: 12px; background: #1e293b; color: #fff; border: 1px solid #475569; border-radius: 8px; outline: none; box-sizing: border-box; font-size: 14px; }
        input:focus { border-color: var(--primary); }
        button { width: 100%; padding: 14px; background: var(--primary); color: white; font-weight: bold; border: none; border-radius: 8px; cursor: pointer; font-size: 16px; transition: 0.3s; margin-top: 10px; }
        button:hover { background: #2563eb; transform: translateY(-1px); }
        button:disabled { background: #475569; cursor: not-allowed; transform: none; }
        #result { margin-top: 30px; padding: 20px; background: #064e3b; border-radius: 12px; display: none; animation: fadeIn 0.5s ease; }
        .url-box { background: #0f172a; padding: 12px; border: 1px solid #059669; border-radius: 6px; margin: 10px 0; font-family: monospace; font-size: 13px; color: #34d399; overflow-x: auto; white-space: nowrap; }
        .stats { display: flex; justify-content: space-between; font-size: 13px; margin-top: 15px; background: rgba(0,0,0,0.2); padding: 8px 12px; border-radius: 6px; }
        .preview-img { max-width: 100%; border-radius: 8px; margin-top: 15px; border: 1px solid #334155; display: block; }
        .notice-text { color: #fbbf24; font-size: 13.5px; margin-top: 15px; font-weight: 600; text-align: center; background: rgba(251, 191, 36, 0.1); padding: 10px; border-radius: 6px; border: 1px solid rgba(251, 191, 36, 0.3); }
        @keyframes fadeIn { from { opacity: 0; transform: translateY(10px); } to { opacity: 1; transform: translateY(0); } }
        .copy-btn { background: #334155; color: #fff; padding: 5px 10px; font-size: 11px; border-radius: 4px; border: none; cursor: pointer; float: right; margin-top: -35px; margin-right: 5px; }
    </style>
</head>
<body>
    <div class="card">
        <h2>🚀 OptiPix Pro</h2>
        <p class="sub">High-performance AI image compression and cloud hosting.</p>
        
        <div class="input-group">
            <label>Remote Image URL</label>
            <input type="text" id="imgUrl" placeholder="https://example.com/image.jpg">
        </div>
        
        <div style="text-align: center; margin-bottom: 20px; color: #475569; font-size: 12px;">— OR —</div>

        <div class="input-group">
            <label>Upload Local Image</label>
            <input type="file" id="imgFile" accept="image/*">
        </div>

        <label>Compression Level</label>
        <select id="level">
            <option value="none">Original (No Compression)</option>
            <option value="extreme" selected>Web-Ready (Max Optimization)</option>
            <option value="medium">Standard (Balanced)</option>
        </select>
        
        <button onclick="processImage()" id="btn">Optimize & Host Image</button>
        
        <div id="result">
            <h3 style="margin:0 0 10px 0; font-size: 18px; color: #34d399;">✅ Processing Complete</h3>
            <p style="margin: 0; color: #a7f3d0; font-size: 13px;">CDN Optimized Link:</p>
            <div class="url-box" id="resLink"></div>
            <button class="copy-btn" onclick="copyUrl()">Copy URL</button>
            
            <div class="stats">
                <span>Original: <span id="oldSize">--</span> KB</span>
                <span style="color: #34d399;">Optimized: <span id="newSize">--</span> KB</span>
            </div>
            
            <div class="notice-text">
                ⏳ Note: The image URL will be live in max 30 seconds.
            </div>
            
            <img id="resImg" class="preview-img" src="" alt="Preview will load soon...">
        </div>
    </div>

    <script>
        async function processImage() {
            const btn = document.getElementById('btn');
            const urlInput = document.getElementById('imgUrl').value;
            const fileInput = document.getElementById('imgFile').files[0];
            const level = document.getElementById('level').value;
            const resDiv = document.getElementById('result');
            
            if(!urlInput && !fileInput) { alert("Please provide an image source!"); return; }
            
            btn.innerText = "⚡ Processing on Cloud...";
            btn.disabled = true;
            resDiv.style.display = "none";
            
            let formData = new FormData();
            formData.append('level', level);
            if(fileInput) formData.append('file', fileInput);
            else formData.append('url', urlInput);

            try {
                const response = await fetch('/upload-poster', { method: 'POST', body: formData });
                const data = await response.json();
                
                if(data.success) {
                    document.getElementById('resLink').innerText = data.url;
                    document.getElementById('oldSize').innerText = data.original_size_kb;
                    document.getElementById('newSize').innerText = data.encrypted_size_kb;
                    
                    document.getElementById('resImg').src = "";
                    
                    // 30 seconds wait timeout (plus 1s buffer)
                    setTimeout(() => {
                        document.getElementById('resImg').src = data.url;
                    }, 31000);

                    resDiv.style.display = "block";
                } else {
                    alert("Processing Error: " + data.error);
                }
            } catch (err) {
                alert("Cloud connection failed!");
            }
            btn.innerText = "Optimize & Host Image";
            btn.disabled = false;
        }

        function copyUrl() {
            const url = document.getElementById('resLink').innerText;
            navigator.clipboard.writeText(url);
            alert("Link copied to clipboard!");
        }
    </script>
</body>
</html>
"""

@app.route('/')
def index():
    return render_template_string(HTML_UI)

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