backend/api/translate.py

import datetime
import os
import requests as py_requests
from flask import Blueprint, request, jsonify, Response, session, current_app
from backend.config import Config
from backend.core.decorators import get_current_user
from backend.core.rate_limit import get_client_ip
from backend.core.security import verify_access_token
from backend.services.translation import (
    get_engine, translate_texts, translate_stream_generator,
    translation_limit_tracker, call_ai_chat_proxy
)

translate_bp = Blueprint("translate", __name__)


def is_vip_request():
    """Check if current request has VIP privileges (session, JWT, or header code)."""
    user = get_current_user()
    if user and user.get("vip_status", 0) == 1:
        return True
    
    # Check headers
    vip_code = request.headers.get("X-VIP-Code", "") or request.headers.get("X-VIP-Key", "")
    if not vip_code:
        # Check request body
        try:
            data = request.json or {}
            vip_code = data.get("vip_key", "") or data.get("vip_code", "")
        except:
            pass
            
    if vip_code and vip_code in Config.VALID_VIP_CODES:
        return True
    return False


def _check_translation_rate_limit(texts):
    """Returns (exceeded: bool, tracker_key, count) for standard users."""
    client_ip = get_client_ip()
    today_str = datetime.date.today().isoformat()
    tracker_key = f"{client_ip}:{today_str}"
    current_count = translation_limit_tracker.get(tracker_key, 0)
    requested_count = len([t for t in texts if t.strip()])
    exceeded = (current_count + requested_count) > 50
    return exceeded, tracker_key, requested_count


@translate_bp.route("/translate", methods=["POST"])
def translate():
    data = request.json
    if not data or "texts" not in data:
        return jsonify({"error": "Missing 'texts' array"}), 400

    texts = data["texts"]

    if not is_vip_request():
        exceeded, tracker_key, count = _check_translation_rate_limit(texts)
        if exceeded:
            return jsonify({
                "error": "Hạn mức dịch máy chủ Standard đã hết (50 đoạn/ngày). Vui lòng nâng cấp tài khoản VIP!"
            }), 403
        translation_limit_tracker[tracker_key] = translation_limit_tracker.get(tracker_key, 0) + count

    mode = data.get("mode")

    # If standalone translation server is running, proxy request to it
    local_translate_url = os.environ.get("LOCAL_TRANSLATE_URL", "")
    if not local_translate_url:
        local_tts_url = os.environ.get("LOCAL_TTS_URL", "")
        if local_tts_url:
            local_translate_url = local_tts_url.split("/v1/")[0]

    if local_translate_url:
        try:
            translate_url = f"{local_translate_url.rstrip('/')}/v1/translate"
            r = py_requests.post(translate_url, json={"texts": texts, "mode": mode}, timeout=(2.0, 60.0))
            if r.status_code == 200:
                return jsonify(r.json())
            current_app.logger.error(f"[Translate Proxy Error]: {r.status_code} - {r.text}")
        except Exception as e:
            current_app.logger.error(f"[Translate Proxy Exception]: {e}")

    translations = translate_texts(texts, mode)
    return jsonify({"translations": translations})


@translate_bp.route("/translate_stream", methods=["POST"])
def translate_stream():
    data = request.json
    if not data or "texts" not in data:
        return jsonify({"error": "Missing 'texts' array"}), 400

    texts = data["texts"]

    if not is_vip_request():
        exceeded, tracker_key, count = _check_translation_rate_limit(texts)
        if exceeded:
            return jsonify({
                "error": "Hạn mức dịch máy chủ Standard đã hết (50 đoạn/ngày). Vui lòng nâng cấp tài khoản VIP!"
            }), 403
        translation_limit_tracker[tracker_key] = translation_limit_tracker.get(tracker_key, 0) + count

    mode = data.get("mode")

    # If standalone translation server is running, proxy stream request to it
    local_translate_url = os.environ.get("LOCAL_TRANSLATE_URL", "")
    if not local_translate_url:
        local_tts_url = os.environ.get("LOCAL_TTS_URL", "")
        if local_tts_url:
            local_translate_url = local_tts_url.split("/v1/")[0]

    if local_translate_url:
        try:
            translate_url = f"{local_translate_url.rstrip('/')}/v1/translate_stream"
            r = py_requests.post(translate_url, json={"texts": texts, "mode": mode}, stream=True, timeout=(2.0, 60.0))
            if r.status_code == 200:
                def stream_forwarder():
                    for chunk in r.iter_content(chunk_size=1024):
                        if chunk:
                            yield chunk
                response = Response(stream_forwarder(), mimetype="text/event-stream")
                response.headers["X-Accel-Buffering"] = "no"
                response.headers["Cache-Control"] = "no-cache, no-transform"
                response.headers["Connection"] = "keep-alive"
                response.headers["X-Content-Type-Options"] = "nosniff"
                return response
            current_app.logger.error(f"[Translate Stream Proxy Error]: {r.status_code} - {r.text}")
        except Exception as e:
            current_app.logger.error(f"[Translate Stream Proxy Exception]: {e}")

    response = Response(translate_stream_generator(texts, mode), mimetype="text/event-stream")
    response.headers["X-Accel-Buffering"] = "no"
    response.headers["Cache-Control"] = "no-cache, no-transform"
    response.headers["Connection"] = "keep-alive"
    response.headers["X-Content-Type-Options"] = "nosniff"
    return response


@translate_bp.route("/api/ai/chat", methods=["POST"])
def ai_chat_proxy():
    if not is_vip_request():
        return jsonify({
            "error": "Chức năng AI chỉ dành cho thành viên VIP. Vui lòng nâng cấp hoặc nhập API Key cá nhân trong cài đặt!"
        }), 403

    data = request.json or {}
    messages = data.get("messages", [])
    model = data.get("model", "gemini-1.5-flash")
    prompt = data.get("prompt", "")

    try:
        text = call_ai_chat_proxy(messages, model, prompt)
        return jsonify({"text": text})
    except ValueError as e:
        return jsonify({"error": str(e)}), 501
    except RuntimeError as e:
        return jsonify({"error": str(e)}), 502
    except Exception as e:
        return jsonify({"error": f"Lỗi xử lý AI Proxy: {str(e)}"}), 500


@translate_bp.route("/iframe_proxy", methods=["GET"])
def iframe_proxy():
    url = request.args.get("url")
    if not url:
        return "Missing 'url' parameter", 400

    try:
        headers = {
            "User-Agent": "Mozilla/5.0 (Linux; Android 10; K) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Mobile Safari/537.36",
            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
            "Accept-Language": "en-US,en;q=0.5"
        }
        
        # Follow redirects
        r = py_requests.get(url, headers=headers, timeout=15, allow_redirects=True)
        
        encoding = r.encoding
        if not encoding or encoding.lower() == 'iso-8859-1':
            encoding = r.apparent_encoding or 'utf-8'
        
        try:
            html = r.content.decode(encoding, errors='ignore')
        except Exception:
            html = r.content.decode('utf-8', errors='ignore')

        # Insert base URL tag so all links/resources load relative to original page
        base_tag = f'<base href="{url}">'
        if "<head>" in html:
            html = html.replace("<head>", f"<head>{base_tag}", 1)
        elif "<HEAD>" in html:
            html = html.replace("<HEAD>", f"<HEAD>{base_tag}", 1)
        else:
            html = base_tag + html

        # Injected proxy handler script
        injected_script = """
        <script>
        (function() {
            const origLog = console.log;
            console.log = function(...args) {
                origLog.apply(console, args);
                const msg = args.join(' ');
                if (msg === '[Translation Complete]') {
                    window.parent.postMessage({ type: 'TRANSLATION_COMPLETE' }, '*');
                } else if (msg.startsWith('[TRANSLATE_REQ]')) {
                    try {
                        const payload = JSON.parse(msg.substring(15));
                        window.parent.postMessage({ type: 'TRANSLATE_REQ', payload: payload }, '*');
                    } catch(e) {
                        origLog.error("[Proxy Iframe] Parse translation req error:", e);
                    }
                }
            };

            window.addEventListener('message', function(event) {
                const data = event.data;
                if (!data || !data.action) return;

                if (data.action === 'INJECT_SCRIPT') {
                    try {
                        const scriptEl = document.createElement('script');
                        scriptEl.textContent = data.script;
                        document.body.appendChild(scriptEl);
                    } catch (err) {
                        origLog.error("[Proxy Iframe] Error injecting script:", err);
                    }
                } else if (data.action === 'translate') {
                    if (window.toggleAutoTranslate) {
                        window.toggleAutoTranslate(data.enabled);
                    }
                } else if (data.action === 'audio') {
                    let res = { title: document.title, text: document.body.innerText };
                    if (window.__TienHiepHelpers) {
                        res = window.__TienHiepHelpers.extractCleanChapterText();
                    }
                    window.parent.postMessage({
                        type: 'AUDIO_TEXT_RES',
                        title: res.title,
                        text: res.text
                    }, '*');
                } else if (data.action === 'scroll') {
                    if (window.__scrollInterval) {
                        clearInterval(window.__scrollInterval);
                        window.__scrollInterval = null;
                    } else {
                        const speed = data.speed || 30;
                        window.__scrollInterval = setInterval(() => {
                            window.scrollBy({top: 1, behavior: 'instant'});
                        }, speed);
                    }
                } else if (data.action === 'next') {
                    if (window.__TienHiepHelpers) {
                        window.__TienHiepHelpers.checkAndTriggerAutoNext(true);
                    }
                } else if (data.action === 'teach_next') {
                    if (window.__TienHiepHelpers) {
                        window.__TienHiepHelpers.startTeachNextMode();
                    }
                } else if (data.action === 'dark_mode') {
                    if (document.documentElement.style.filter.includes('invert(1)')) {
                        document.documentElement.style.filter = '';
                        document.documentElement.style.backgroundColor = '';
                    } else {
                        document.documentElement.style.filter = 'invert(1) hue-rotate(180deg) brightness(0.9) contrast(1.1)';
                        document.documentElement.style.backgroundColor = '#121212';
                    }
                } else if (data.action === 'clean_ads') {
                    const ads = document.querySelectorAll('iframe, .ad, .ads, [id*="ad"], [class*="ad"], .banner, .popup, ins');
                    ads.forEach(ad => ad.remove());
                } else if (data.action === 'force_translate') {
                    if (window.__autoTranslateObserver) {
                         const allNodes = window.__TienHiepHelpers ? window.__TienHiepHelpers.collectTextNodes(document.body) : [];
                         if (allNodes.length > 0) window.__translateQueue = allNodes;
                         if (window.__processTranslationQueue) window.__processTranslationQueue();
                    }
                } else if (data.action === 'SET_TTS_PLAYING') {
                    window.isTtsPlaying = data.playing;
                } else if (data.action === 'highlight') {
                    const targetText = data.sentenceText ? data.sentenceText.trim() : '';
                    if (targetText) {
                        function findTextNode(root, text) {
                          const walker = document.createTreeWalker(root, NodeFilter.SHOW_TEXT, null, false);
                          let node;
                          while (node = walker.nextNode()) {
                            const val = node.nodeValue || '';
                            if (val.includes(text)) {
                              return { node, startIdx: val.indexOf(text) };
                            }
                          }
                          if (text.length > 15) {
                            const prefix = text.substring(0, Math.floor(text.length * 0.65));
                            const walker2 = document.createTreeWalker(root, NodeFilter.SHOW_TEXT, null, false);
                            while (node = walker2.nextNode()) {
                              const val = node.nodeValue || '';
                              if (val.includes(prefix)) {
                                return { node, startIdx: val.indexOf(prefix) };
                              }
                            }
                          }
                          return null;
                        }

                        let attempts = 0;
                        function tryHighlight() {
                          const oldHighlight = document.getElementById('tienhiep-active-highlight');
                          const result = findTextNode(document.body, targetText);
                          
                          if (result) {
                            if (oldHighlight) {
                              const parent = oldHighlight.parentNode;
                              if (parent) {
                                const textNode = document.createTextNode(oldHighlight.textContent);
                                parent.replaceChild(textNode, oldHighlight);
                                parent.normalize();
                              }
                            }

                            const { node, startIdx } = result;
                            const parent = node.parentNode;
                            if (parent && parent.nodeName !== 'SCRIPT' && parent.nodeName !== 'STYLE') {
                              const textVal = node.nodeValue;
                              const matchLen = Math.min(targetText.length, textVal.length - startIdx);
                              
                              const beforeText = textVal.substring(0, startIdx);
                              const matchedText = textVal.substring(startIdx, startIdx + matchLen);
                              const afterText = textVal.substring(startIdx + matchLen);

                              const fragment = document.createDocumentFragment();
                              if (beforeText) fragment.appendChild(document.createTextNode(beforeText));
                              
                              const span = document.createElement('span');
                              span.id = 'tienhiep-active-highlight';
                              span.style.backgroundColor = 'rgba(139, 92, 246, 0.25)';
                              span.style.color = '#c084fc';
                              span.style.borderBottom = '2px solid #a855f7';
                              span.style.padding = '1px 3px';
                              span.style.borderRadius = '3px';
                              span.style.transition = 'all 0.3s ease';
                              span.textContent = matchedText;
                              fragment.appendChild(span);

                              if (afterText) fragment.appendChild(document.createTextNode(afterText));

                              parent.replaceChild(fragment, node);
                              span.scrollIntoView({ behavior: 'smooth', block: 'center' });
                            }
                          } else if (attempts < 6) {
                            attempts++;
                            setTimeout(tryHighlight, 600);
                          }
                        }

                        tryHighlight();
                    }
                }
            });

            document.addEventListener('click', function(e) {
                const a = e.target.closest('a');
                if (a && a.href) {
                    if (a.href.startsWith('http://') || a.href.startsWith('https://')) {
                        e.preventDefault();
                        window.parent.postMessage({
                            type: 'NAVIGATE_REQ',
                            url: a.href
                        }, '*');
                    }
                }
            }, true);

            window.addEventListener('DOMContentLoaded', function() {
                window.parent.postMessage({ type: 'IFRAME_READY', url: window.location.href }, '*');
            });
            if (document.readyState === 'interactive' || document.readyState === 'complete') {
                window.parent.postMessage({ type: 'IFRAME_READY', url: window.location.href }, '*');
            }
        })();
        </script>
        """

        if "</body>" in html:
            html = html.replace("</body>", f"{injected_script}</body>", 1)
        elif "</BODY>" in html:
            html = html.replace("</BODY>", f"{injected_script}</BODY>", 1)
        else:
            html = html + injected_script

        response = Response(html, mimetype="text/html")
        response.headers.pop("X-Frame-Options", None)
        response.headers.pop("Content-Security-Policy", None)
        response.headers["Access-Control-Allow-Origin"] = "*"
        return response
    except Exception as e:
        return f"Proxy error: {str(e)}", 500