app.py

# app.py
import os
import json
import time
from flask import Flask, request, Response, jsonify, render_template, redirect, url_for
from google import genai
from google.genai import types
from pydantic import BaseModel, Field
import psycopg2
import psycopg2.pool

app = Flask(__name__, template_folder=".")

# --- Database Setup (Neon PostgreSQL) ---
DATABASE_URL = os.environ.get("DATABASE_URL")
if not DATABASE_URL:
    raise RuntimeError("DATABASE_URL environment variable is required (e.g. from Neon.tech)")

# Connection pool for better performance
try:
    db_pool = psycopg2.pool.SimpleConnectionPool(1, 10, DATABASE_URL)
    print("Database connection pool created")
except Exception as e:
    print(f"Failed to create database pool: {e}")
    raise

def get_db_connection():
    """Get a connection from the pool."""
    return db_pool.getconn()

def release_db_connection(conn):
    """Return a connection to the pool."""
    db_pool.putconn(conn)

def init_db():
    """Create tables if they don't exist."""
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            # Menu table
            cur.execute("""
                CREATE TABLE IF NOT EXISTS menu (
                    id SERIAL PRIMARY KEY,
                    name TEXT NOT NULL,
                    description TEXT DEFAULT '',
                    price TEXT DEFAULT '',
                    available BOOLEAN DEFAULT TRUE,
                    created_at TIMESTAMP DEFAULT NOW()
                );
            """)
            # Orders table
            cur.execute("""
                CREATE TABLE IF NOT EXISTS orders (
                    id SERIAL PRIMARY KEY,
                    table_number TEXT NOT NULL,
                    items JSONB NOT NULL,
                    status TEXT DEFAULT 'pending',
                    timestamp DOUBLE PRECISION DEFAULT 0,
                    created_at TIMESTAMP DEFAULT NOW()
                );
            """)
        conn.commit()
    except Exception as e:
        print(f"Database initialization error: {e}")
        conn.rollback()
    finally:
        release_db_connection(conn)

# Initialize database on startup
init_db()

# --- In-memory caches (thread-safe session data) ---
table_chat_histories = {}   # conversation history per table
global_drafts = {}          # draft orders per table
active_tasks = {}           # stream cancellation tokens

# --- Pydantic Models ---
class OrderItem(BaseModel):
    name: str = Field(description="نام دقیق و کامل آیتم از منوی فعال به زبان فارسی (مثال: کیک شکلاتی خیس)")
    quantity: int = Field(description="تعداد درخواستی مشتری از این آیتم")

# --- Database helper functions ---
def load_menu():
    """Load all menu items from database."""
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute("SELECT id, name, description, price, available FROM menu ORDER BY id;")
            rows = cur.fetchall()
            menu = []
            for row in rows:
                menu.append({
                    "id": row[0],
                    "name": row[1],
                    "description": row[2],
                    "price": row[3],
                    "available": row[4]
                })
            return menu
    except Exception as e:
        print(f"load_menu error: {e}")
        return []
    finally:
        release_db_connection(conn)

def save_menu(menu_data):
    """Replace the entire menu list (used when admin publishes a new menu)."""
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            # Clear existing menu
            cur.execute("DELETE FROM menu;")
            # Insert new items
            for item in menu_data:
                cur.execute(
                    "INSERT INTO menu (name, description, price, available) VALUES (%s, %s, %s, %s)",
                    (item.get("name", ""), item.get("description", ""),
                     item.get("price", ""), item.get("available", True))
                )
        conn.commit()
        return True
    except Exception as e:
        print(f"save_menu error: {e}")
        conn.rollback()
        return False
    finally:
        release_db_connection(conn)

def load_orders():
    """Load all orders from database."""
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute("SELECT id, table_number, items, status, timestamp FROM orders ORDER BY id;")
            rows = cur.fetchall()
            orders = []
            for row in rows:
                orders.append({
                    "id": row[0],
                    "table": row[1],
                    "items": row[2],          # already JSON
                    "status": row[3],
                    "timestamp": row[4]
                })
            return orders
    except Exception as e:
        print(f"load_orders error: {e}")
        return []
    finally:
        release_db_connection(conn)

# --- Tool functions (will be called by Gemini) ---
def get_menu() -> str:
    """
    Returns the current menu items. Use this to see what drinks, foods, or sweets are available in the cafe.
    """
    menu = load_menu()
    available_items = [item for item in menu if item.get("available", True)]
    if not available_items:
        return "در حال حاضر هیچ منویی آپلود نشده است و منوی فعال موجود نیست."
    # Return only needed fields for AI (name, description, price)
    simplified = [{"name": i["name"], "description": i["description"], "price": i["price"]} for i in available_items]
    return json.dumps(simplified, ensure_ascii=False)

def prepare_order_draft(table_number: str, items: list[OrderItem]) -> str:
    """
    Prepares a draft order list for the customer to confirm on their screen.
    Only call this after verifying that every item exists in the menu (using get_menu).
    """
    table_number_str = str(table_number)
    items_list = []
    for item in items:
        if isinstance(item, dict):
            items_list.append({
                "name": item.get("name", ""),
                "quantity": item.get("quantity", 1)
            })
        elif hasattr(item, "name"):
            items_list.append({
                "name": item.name,
                "quantity": item.quantity
            })
        else:
            # fallback
            items_list.append({
                "name": str(item),
                "quantity": 1
            })
    global_drafts[table_number_str] = {
        "table": table_number_str,
        "items": items_list,
        "timestamp": time.time()
    }
    return f"من پیش‌نویس سفارش شما رو برای میز {table_number_str} آماده کردم."

def get_full_menu() -> str:
    """
    Returns the complete menu of the cafe, including both available and unavailable items.
    """
    menu = load_menu()
    if not menu:
        return "منوی کافه خالی است."
    return json.dumps(menu, ensure_ascii=False)

def set_item_availability(item_name: str, available: bool) -> str:
    """
    Updates whether a menu item is available or sold out.
    """
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            # Find item by name (case-insensitive)
            cur.execute("SELECT id, name FROM menu WHERE LOWER(name) = %s;", (item_name.strip().lower(),))
            row = cur.fetchone()
            if row:
                cur.execute("UPDATE menu SET available = %s WHERE id = %s;", (available, row[0]))
                conn.commit()
                status = "موجود" if available else "ناموجود"
                return f"وضعیت آیتم '{row[1]}' به موفقیت به '{status}' تغییر یافت."
            else:
                return f"آیتم '{item_name}' در منوی کافه یافت نشد."
    except Exception as e:
        print(f"set_item_availability error: {e}")
        conn.rollback()
        return f"خطا در بروزرسانی آیتم: {str(e)}"
    finally:
        release_db_connection(conn)

# --- NEW: Full CRUD tools for admin chat ---
def get_menu_items() -> str:
    """
    Returns the complete menu with IDs, names, descriptions, prices, and availability statuses.
    """
    menu = load_menu()
    return json.dumps(menu, ensure_ascii=False)

def add_menu_item(name: str, description: str = "", price: str = "", available: bool = True) -> str:
    """
    Add a new item to the cafe menu. Returns confirmation message.
    """
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                "INSERT INTO menu (name, description, price, available) VALUES (%s, %s, %s, %s) RETURNING id;",
                (name.strip(), description.strip(), price.strip(), available)
            )
            new_id = cur.fetchone()[0]
        conn.commit()
        return f"آیتم '{name}' با شناسه {new_id} با موفقیت اضافه شد."
    except Exception as e:
        conn.rollback()
        return f"خطا در افزودن آیتم: {str(e)}"
    finally:
        release_db_connection(conn)

def update_menu_item(item_id: int, name: str = None, description: str = None, price: str = None, available: bool = None) -> str:
    """
    Update an existing menu item by ID. Only provided fields are updated.
    """
    conn = get_db_connection()
    try:
        updates = {}
        if name is not None:
            updates["name"] = name.strip()
        if description is not None:
            updates["description"] = description.strip()
        if price is not None:
            updates["price"] = price.strip()
        if available is not None:
            updates["available"] = available
        if not updates:
            return "هیچ فیلدی برای بروزرسانی ارسال نشده است."
        with conn.cursor() as cur:
            set_clause = ", ".join([f"{field} = %s" for field in updates])
            values = list(updates.values()) + [item_id]
            cur.execute(f"UPDATE menu SET {set_clause} WHERE id = %s;", values)
            conn.commit()
            if cur.rowcount == 0:
                return f"آیتم با شناسه {item_id} یافت نشد."
            return f"آیتم {item_id} با موفقیت بروزرسانی شد."
    except Exception as e:
        conn.rollback()
        return f"خطا در بروزرسانی آیتم: {str(e)}"
    finally:
        release_db_connection(conn)

def delete_menu_item(item_id: int) -> str:
    """
    Delete a menu item by ID.
    """
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute("DELETE FROM menu WHERE id = %s;", (item_id,))
            conn.commit()
            if cur.rowcount == 0:
                return f"آیتم با شناسه {item_id} یافت نشد."
            return f"آیتم {item_id} با موفقیت حذف شد."
    except Exception as e:
        conn.rollback()
        return f"خطا در حذف آیتم: {str(e)}"
    finally:
        release_db_connection(conn)

# --- System Instructions ---
CUSTOMER_SYSTEM_INSTRUCTION = """Your name is Nila, a classic, warm, and highly professional AI cafe assistant for "Cafe AI". You were developed by "Nastaran Data Algorithm".

You are currently talking to a customer at a specific table. Your responsibilities:
1. Greet warmly and take orders in Persian (you can understand English but always reply in Persian unless the customer insists).
2. **Before preparing any order draft**, you MUST check the current menu by calling the tool **get_menu()** to see what items are available.
3. If the customer asks for an item that is NOT in the menu (exact or similar name), politely say: "متأسفانه این آیتم در منوی امروز ما موجود نیست. می‌توانم پیشنهادهای دیگری از منو به شما بدهم؟" and then suggest 2-3 similar or popular items from the current menu. Never prepare a draft with items not in the menu.
4. Once you have verified all requested items exist, use **prepare_order_draft** to create the order.
5. After preparing the draft, tell the customer they can review and confirm on their screen.
6. If the menu is empty (no items returned), tell the customer that the cafe is currently not taking orders and to wait for the staff.

Always be friendly, speak in Persian, and use the tools to provide accurate information.
"""

ADMIN_SYSTEM_INSTRUCTION = """Your name is Nila, the smart cafe administrator assistant for "Cafe AI". You manage the cafe menu in real time.
You can view, add, update, delete menu items, and change availability using the following tools:
- **get_menu_items()**: fetch the full menu with IDs, names, descriptions, prices, and availability.
- **add_menu_item(name, description, price, available)**: add a new item (description and price optional, available is True by default).
- **update_menu_item(item_id, name, description, price, available)**: update any fields of an existing item by its ID (only pass the fields you want to change).
- **delete_menu_item(item_id)**: permanently remove an item.
- **set_item_availability(item_name, available)**: quickly toggle availability of an item by its exact name.

When the user asks to see the menu, use get_menu_items and show a nicely formatted list.
When asked to add an item, ask for the necessary details and then call add_menu_item.
For modifications, first get the menu to find the item ID, then use update_menu_item.
For removals, use delete_menu_item with the correct ID.
Always confirm every change clearly. Communicate primarily in Persian (Farsi), but you can understand English commands as well.
Be polite, precise, and efficient."""

# --- Routes ---
@app.route('/')
def home():
    return render_template("customer.html", table_number=None)

@app.route('/customer/<table_number>')
def customer_chat_page(table_number):
    return render_template("customer.html", table_number=table_number)

@app.route('/dashboard')
def admin_dashboard_page():
    return render_template("cafe.html")

# --- Customer API (with function calling loop) ---
@app.route('/api/customer/chat_stream/<table_number>', methods=['POST'])
def api_customer_chat_stream(table_number):
    api_key = os.environ.get("NILLA_CUSTOMER") or os.environ.get("GEMINI_API_KEY")
    if not api_key:
        return jsonify({"error": "API key not found"}), 500

    data = request.json or {}
    user_message = data.get("message")
    if not user_message:
        return jsonify({"error": "Empty message"}), 400

    if table_number not in table_chat_histories:
        table_chat_histories[table_number] = []

    global_drafts.pop(table_number, None)
    table_chat_histories[table_number].append({"role": "user", "content": user_message})
    active_tasks[table_number] = True

    def generate_chunks():
        try:
            # Build conversation history for the model
            contents = []
            for msg in table_chat_histories[table_number][:-1]:
                role = "user" if msg["role"] == "user" else "model"
                contents.append(
                    types.Content(
                        role=role,
                        parts=[types.Part.from_text(text=msg["content"])]
                    )
                )

            client = genai.Client(api_key=api_key)
            full_generated_text = ""

            dynamic_instruction = CUSTOMER_SYSTEM_INSTRUCTION + f"\nYou are currently serving table: {table_number}."

            config = types.GenerateContentConfig(
                system_instruction=dynamic_instruction,
                tools=[get_menu, prepare_order_draft],
                thinking_config=types.ThinkingConfig(thinking_budget=0)
            )

            # Create a chat session
            chat = client.chats.create(
                model="gemini-3.1-flash-lite",
                history=contents,
                config=config
            )

            # Send user message and handle any function calls
            response = chat.send_message(user_message)

            # Loop to process function calls
            max_iterations = 10
            while max_iterations > 0 and active_tasks.get(table_number, True):
                function_calls = []
                if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
                    for part in response.candidates[0].content.parts:
                        if hasattr(part, "function_call") and part.function_call:
                            function_calls.append(part.function_call)

                if not function_calls:
                    # No more function calls, extract the text
                    if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
                        text_parts = [part.text for part in response.candidates[0].content.parts if hasattr(part, "text") and part.text]
                        full_generated_text = "".join(text_parts)
                    if not full_generated_text:
                        try:
                            full_generated_text = response.text or ""
                        except (ValueError, AttributeError):
                            full_generated_text = "متأسفانه مشکلی پیش آمده، لطفاً دوباره تلاش کنید."
                    break

                # Execute function calls
                function_response_parts = []
                for fc in function_calls:
                    func_name = fc.name
                    args = dict(fc.args) if hasattr(fc, "args") else {}
                    print(f"Customer function call: {func_name}({args})")
                    try:
                        if func_name == "get_menu":
                            result = get_menu()
                        elif func_name == "prepare_order_draft":
                            result = prepare_order_draft(**args)
                        else:
                            result = f"Unknown function: {func_name}"
                    except Exception as e:
                        result = f"Error executing {func_name}: {str(e)}"
                    function_response_parts.append(
                        types.Part.from_function_response(name=func_name, response={"result": result})
                    )

                # Send the results back to the model
                response = chat.send_message(
                    types.Content(role="user", parts=function_response_parts)
                )
                max_iterations -= 1

            # If no text was generated after all iterations
            if not full_generated_text:
                full_generated_text = "پاسخی دریافت نشد، لطفاً دوباره تلاش کنید."

            # Save the model's final text to history
            table_chat_histories[table_number].append({"role": "model", "content": full_generated_text})

            # Stream the text chunk by chunk (simulate streaming)
            chunk_size = 4
            for i in range(0, len(full_generated_text), chunk_size):
                if not active_tasks.get(table_number, True):
                    break
                chunk = full_generated_text[i:i+chunk_size]
                yield f"data: {json.dumps({'type': 'text', 'content': chunk}, ensure_ascii=False)}\n\n"
                time.sleep(0.015)

            # After streaming text, check if there's a draft for this table
            draft = global_drafts.get(table_number)
            if draft:
                yield f"data: {json.dumps({'type': 'draft', 'items': draft['items']}, ensure_ascii=False)}\n\n"

        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'content': str(e)}, ensure_ascii=False)}\n\n"
        finally:
            active_tasks.pop(table_number, None)

    return Response(generate_chunks(), mimetype='text/event-stream')

@app.route('/api/customer/stop/<table_number>', methods=['POST'])
def api_customer_stop(table_number):
    active_tasks[table_number] = False
    return jsonify({"success": True})

@app.route('/api/confirm_order', methods=['POST'])
def api_confirm_order():
    data = request.json or {}
    table_number = data.get("table_number")
    items = data.get("items")

    if not table_number:
        return jsonify({"error": "شماره میز مشخص نیست."}), 400

    if not items:
        draft = global_drafts.get(table_number)
        if not draft:
            return jsonify({"error": "اطلاعات پیش‌نویس یافت نشد."}), 404
        items = draft.get("items")

    # Insert order into database
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                "INSERT INTO orders (table_number, items, status, timestamp) VALUES (%s, %s, 'pending', %s) RETURNING id;",
                (table_number, json.dumps(items), time.time())
            )
            new_id = cur.fetchone()[0]
        conn.commit()

        # --- MODIFICATION: Do NOT clear the session so the customer can continue chatting ---
        # The draft is already used, but the chat history remains intact.
        # Remove the draft for this table (a new draft will be created for subsequent orders).
        global_drafts.pop(table_number, None)

        return jsonify({"success": True, "message": "سفارش ثبت شد. می‌توانید همچنان با من گفتگو کنید."})
    except Exception as e:
        conn.rollback()
        return jsonify({"error": f"خطا در ثبت سفارش: {str(e)}"}), 500
    finally:
        release_db_connection(conn)

# --- Admin API ---
@app.route('/api/admin/orders', methods=['GET'])
def api_admin_orders():
    orders = load_orders()
    active_orders = [o for o in orders if o.get("status") == "pending"]
    return jsonify(active_orders)

@app.route('/api/admin/complete_order', methods=['POST'])
def api_admin_complete_order():
    data = request.json or {}
    order_id = data.get("order_id")
    if not order_id:
        return jsonify({"error": "شناسه سفارش نامعتبر"}), 400

    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute("UPDATE orders SET status = 'completed' WHERE id = %s;", (order_id,))
            conn.commit()
            if cur.rowcount == 0:
                return jsonify({"error": "سفارش پیدا نشد."}), 404
            return jsonify({"success": True, "message": "سفارش با موفقیت برداشته شد."})
    except Exception as e:
        conn.rollback()
        return jsonify({"error": str(e)}), 500
    finally:
        release_db_connection(conn)

@app.route('/api/admin/chat', methods=['POST'])
def api_admin_chat():
    api_key = os.environ.get("NILLA_CAFE") or os.environ.get("GEMINI_API_KEY")
    if not api_key:
        return jsonify({"error": "API key not found"}), 500

    data = request.json or {}
    messages = data.get("messages", [])
    if not messages:
        return jsonify({"error": "Empty history"}), 400

    try:
        client = genai.Client(api_key=api_key)
        contents = []
        for msg in messages[:-1]:
            role = "user" if msg["role"] == "user" else "model"
            contents.append(
                types.Content(
                    role=role,
                    parts=[types.Part.from_text(text=msg["content"])]
                )
            )

        config = types.GenerateContentConfig(
            system_instruction=ADMIN_SYSTEM_INSTRUCTION,
            tools=[get_menu_items, add_menu_item, update_menu_item, delete_menu_item, set_item_availability],
            thinking_config=types.ThinkingConfig(thinking_budget=0)
        )

        last_user_message = messages[-1]["content"]

        chat = client.chats.create(
            model="gemini-3.1-flash-lite",
            history=contents,
            config=config
        )

        response = chat.send_message(last_user_message)

        # Handle function calls loop
        max_iterations = 10
        while max_iterations > 0:
            function_calls = []
            if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
                for part in response.candidates[0].content.parts:
                    if hasattr(part, "function_call") and part.function_call:
                        function_calls.append(part.function_call)
            if not function_calls:
                break

            function_response_parts = []
            for fc in function_calls:
                func_name = fc.name
                args = dict(fc.args) if hasattr(fc, "args") else {}
                print(f"Admin function call: {func_name}({args})")
                try:
                    if func_name == "get_menu_items":
                        result = get_menu_items()
                    elif func_name == "add_menu_item":
                        result = add_menu_item(**args)
                    elif func_name == "update_menu_item":
                        result = update_menu_item(**args)
                    elif func_name == "delete_menu_item":
                        result = delete_menu_item(**args)
                    elif func_name == "set_item_availability":
                        result = set_item_availability(**args)
                    else:
                        result = f"Unknown function: {func_name}"
                except Exception as e:
                    result = f"Error executing {func_name}: {str(e)}"
                function_response_parts.append(
                    types.Part.from_function_response(name=func_name, response={"result": result})
                )
            response = chat.send_message(
                types.Content(role="user", parts=function_response_parts)
            )
            max_iterations -= 1

        # Extract final text
        admin_response_text = ""
        if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
            parts = response.candidates[0].content.parts
            text_parts = [part.text for part in parts if hasattr(part, "text") and part.text]
            if text_parts:
                admin_response_text = "".join(text_parts)
        if not admin_response_text:
            try:
                admin_response_text = response.text or ""
            except (ValueError, AttributeError):
                admin_response_text = "تغییرات مورد نظر شما با موفقیت اعمال شد."

        return jsonify({"success": True, "response": admin_response_text})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# --- Admin Menu CRUD (manual REST endpoints) ---
@app.route('/api/admin/menu', methods=['GET'])
def api_admin_get_menu():
    menu = load_menu()
    return jsonify(menu)

@app.route('/api/admin/menu/item', methods=['POST'])
def api_admin_add_menu_item():
    data = request.json or {}
    name = data.get("name", "").strip()
    if not name:
        return jsonify({"error": "نام آیتم الزامی است."}), 400
    description = data.get("description", "")
    price = data.get("price", "")
    available = data.get("available", True)

    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                "INSERT INTO menu (name, description, price, available) VALUES (%s, %s, %s, %s) RETURNING id;",
                (name, description, price, available)
            )
            new_id = cur.fetchone()[0]
        conn.commit()
        return jsonify({"success": True, "id": new_id, "message": "آیتم با موفقیت افزوده شد."})
    except Exception as e:
        conn.rollback()
        return jsonify({"error": str(e)}), 500
    finally:
        release_db_connection(conn)

@app.route('/api/admin/menu/item/<int:item_id>', methods=['PUT'])
def api_admin_update_menu_item(item_id):
    data = request.json or {}
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            allowed_fields = ["name", "description", "price", "available"]
            updates = []
            values = []
            for field in allowed_fields:
                if field in data:
                    updates.append(f"{field} = %s")
                    values.append(data[field])
            if not updates:
                return jsonify({"error": "هیچ فیلدی برای بروزرسانی ارسال نشده."}), 400
            values.append(item_id)
            query = f"UPDATE menu SET {', '.join(updates)} WHERE id = %s;"
            cur.execute(query, values)
            conn.commit()
            if cur.rowcount == 0:
                return jsonify({"error": "آیتم پیدا نشد."}), 404
            return jsonify({"success": True, "message": "آیتم بروزرسانی شد."})
    except Exception as e:
        conn.rollback()
        return jsonify({"error": str(e)}), 500
    finally:
        release_db_connection(conn)

@app.route('/api/admin/menu/item/<int:item_id>', methods=['DELETE'])
def api_admin_delete_menu_item(item_id):
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute("DELETE FROM menu WHERE id = %s;", (item_id,))
            conn.commit()
            if cur.rowcount == 0:
                return jsonify({"error": "آیتم پیدا نشد."}), 404
            return jsonify({"success": True, "message": "آیتم حذف شد."})
    except Exception as e:
        conn.rollback()
        return jsonify({"error": str(e)}), 500
    finally:
        release_db_connection(conn)

# --- Legacy extraction & save (still used for bulk upload) ---
@app.route('/api/admin/extract_menu', methods=['POST'])
def api_extract_menu():
    api_key = os.environ.get("NILLA_CAFE") or os.environ.get("GEMINI_API_KEY")
    if not api_key:
        return jsonify({"error": "API key not found"}), 500

    if "file" not in request.files:
        return jsonify({"error": "File missing"}), 400

    file = request.files["file"]
    if file.filename == "":
        return jsonify({"error": "Empty filename"}), 400

    try:
        file_bytes = file.read()
        mime_type = file.mimetype or "image/jpeg"

        prompt = """
        Analyze this menu image or file. Extract all menu items, descriptions, and prices.
        Return ONLY a raw JSON array of objects without markdown formatting.
        Format:
        [
          {"name": "اسپرسو", "description": "قهوه غلیظ خالص", "price": "۵۰,۰۰۰"},
          {"name": "کیک ساده", "description": "شیرینی خانگی لذیذ", "price": "۴۰,۰۰۰"}
        ]
        """

        client = genai.Client(api_key=api_key)
        file_part = types.Part.from_bytes(data=file_bytes, mime_type=mime_type)

        response = client.models.generate_content(
            model="gemini-3.1-flash-lite",
            contents=[file_part, prompt],
            config=types.GenerateContentConfig(
                response_mime_type="application/json"
            )
        )

        extracted_menu = json.loads(response.text)
        for item in extracted_menu:
            item["available"] = True

        return jsonify({"success": True, "extracted_menu": extracted_menu})
    except Exception as e:
        return jsonify({"error": f"خطا در استخراج: {str(e)}"}), 500

@app.route('/api/admin/save_menu', methods=['POST'])
def api_admin_save_menu():
    """Replace entire menu (used after extraction/manual editing)."""
    data = request.json or {}
    menu_list = data.get("menu")
    if not isinstance(menu_list, list):
        return jsonify({"error": "فرمت داده ارسال شده نادرست است."}), 400

    if save_menu(menu_list):
        return jsonify({"success": True, "message": "منوی فعال با موفقیت ثبت شد."})
    else:
        return jsonify({"error": "خطا در ذخیره‌سازی منو."}), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=True)