app.py

# -*- coding: utf-8 -*-
"""
Pumpkin AI Flask Backend - HuggingFace Spaces Version
Serves both API endpoints and frontend static files
"""

from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
import sys
import os
import io
from pathlib import Path
import base64

# Fix Windows console encoding issues
if sys.platform == 'win32':
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

# Import the existing pumpkin logic
from pumpkin_code import (
    load_prompts,
    load_hints_simple,
    load_rescue_lines,
    is_target_lyric,
    classify_input,
    generate_text_reply,
    handle_text_turn,
    handle_image_turn,
    generate_image_data_url,
    save_data_url_to_file,
    ConversationHistory
)

# Initialize Flask app with static file serving
static_folder = Path(__file__).parent / "pumpkin-ai-frontend" / "build"
app = Flask(__name__, static_folder=str(static_folder), static_url_path='')

# Enable CORS for frontend access
CORS(app, origins=[
    "http://localhost:3000",  # Local development
    "http://localhost:3001",
    "http://10.14.0.2:3000",  # Network address
    "https://*.hf.space",  # Hugging Face Spaces
    "https://huggingface.co"  # Hugging Face
])

# Load prompts and data at startup
PROMPT_DIR = Path(__file__).parent / "Prompts"

try:
    classify_prompt, text_reply_prompt, image_reply_prompt, hints_raw = load_prompts(PROMPT_DIR)
    hints_path = PROMPT_DIR / "hints.txt"
    rescue_path = PROMPT_DIR / "rescue_line.txt"
    hints_pool = load_hints_simple(hints_path, limit=10)
    rescue_lines = load_rescue_lines(rescue_path)

    print(f"[OK] Loaded {len(hints_pool)} hints, {len(rescue_lines)} rescue lines")
except Exception as e:
    print(f"[ERROR] Error loading prompts: {e}")
    classify_prompt = text_reply_prompt = image_reply_prompt = ""
    hints_pool = []
    rescue_lines = []

# Store conversation history per user (keyed by IP address)
conversation_histories = {}

def get_conversation_history(user_identifier: str) -> ConversationHistory:
    """Get or create conversation history for a user."""
    if user_identifier not in conversation_histories:
        conversation_histories[user_identifier] = ConversationHistory(max_history=5)
        print(f"[HISTORY] Created new conversation history for {user_identifier}")
    return conversation_histories[user_identifier]

def strip_rescue_and_hint(response: str) -> str:
    """
    Remove rescue line and hint from response for clean history storage.
    This prevents double rescue lines when threading conversations.
    """
    import re
    # Find the separator pattern (one or more equals signs between newlines)
    separator_pattern = r'\n=+\n'
    parts = re.split(separator_pattern, response)
    if len(parts) > 1:
        # Return only the first part (before first separator)
        return parts[0].strip()
    return response

# Serve React frontend
@app.route('/', methods=['GET'])
def serve_frontend():
    """Serve the React frontend index.html"""
    return send_from_directory(str(static_folder), 'index.html')

@app.route('/<path:path>', methods=['GET'])
def serve_static(path):
    """Serve static files (JS, CSS, images)"""
    if static_folder.joinpath(path).exists():
        return send_from_directory(str(static_folder), path)
    else:
        # If file not found, serve index.html (for React Router)
        return send_from_directory(str(static_folder), 'index.html')

@app.route('/api/health', methods=['GET'])
def health():
    """Health check endpoint"""
    return jsonify({
        "status": "healthy",
        "hints_loaded": len(hints_pool),
        "rescue_lines_loaded": len(rescue_lines)
    })

@app.route('/api/chat', methods=['POST'])
def chat():
    """
    Main chat endpoint
    Expected JSON:
    {
        "input": "user message",
        "type": "text" or "image",
        "imageData": "base64 image data" (optional, for image type)
    }

    Returns JSON:
    {
        "response": "AI response or data URL",
        "is_end": true/false,
        "is_image": true/false
    }
    """
    try:
        data = request.get_json()

        if not data or 'input' not in data:
            return jsonify({
                "error": "Missing 'input' field",
                "response": "[WARNING] Please provide input content",
                "is_end": False
            }), 400

        user_input = data['input']
        input_type = data.get('type', 'text')

        # Safe print - avoid encoding errors
        try:
            print(f"[CHAT] Received ({input_type}): {user_input[:50]}...")
        except:
            print(f"[CHAT] Received ({input_type}): [content with special characters]")

        # Check if it's the winning spell
        if is_target_lyric(user_input):
            return jsonify({
                "response": """公主！你終於出現喇！
我有南瓜強迫症，係因為畀一隻南瓜精靈上咗身！南瓜精靈終日幻想自己係灰姑娘故事入面嘅一架南瓜車，只有等到變成公主後嘅灰姑娘出現，佢先肯離開！你唱出咗南瓜界經典金曲《南瓜車》嘅歌詞，超渡咗南瓜精靈，解救咗我，我要好好多謝你！而家就去接收我畀你嘅小禮物啦！Happy Halloween！""",
                "is_end": True
            })

        # Handle text input
        try:
            # Check if user uploaded an image data
            uploaded_image_data = data.get('imageData', None)

            # Skip classification if type is explicitly 'image' or imageData is present
            if input_type == 'image' or uploaded_image_data:
                input_classification = 'image'
                print(f"[IMAGE] Image request detected (type={input_type}, has_imageData={bool(uploaded_image_data)})")
            else:
                # Classify the input to determine if it's an image request
                print(f"[CLASSIFY] Classifying input: {user_input[:50]}...")
                input_classification = classify_input(user_input, classify_prompt)
                print(f"[CLASSIFY] Classification result: {input_classification}")

            # Handle image generation requests
            if input_classification == 'image':
                print(f"[IMAGE] Processing image request: {user_input[:50]}...")

                # Get conversation history for this user
                user_id = request.remote_addr
                conv_history = get_conversation_history(user_id)

                # Log uploaded image data
                if uploaded_image_data:
                    print(f"[IMAGE] User uploaded image data (length: {len(uploaded_image_data)} chars)")

                # Use the updated handle_image_turn that includes rescue lines + hints + analysis + threading
                # Now returns tuple: (text_response, data_url_or_none)
                output_path = "temp_pumpkin_image.png"
                result, data_url = handle_image_turn(
                    user_input,
                    image_reply_prompt,
                    output_path,
                    hints_pool,
                    rescue_lines,
                    uploaded_image_data,
                    conv_history  # Pass conversation history for threading
                )

                # Store this turn in conversation history (simplified description)
                conv_history.add_turn(user_input, "🎃 [生成了南瓜主題圖片]")

                # Check if we got a valid data URL
                if data_url:
                    print(f"[IMAGE] Successfully generated image with data URL (length: {len(data_url)} chars)")
                    print(f"[IMAGE] Data URL preview: {data_url[:100]}...")
                    print(f"[IMAGE] Returning JSON with imageData field")

                    # Return the data URL with the text response that includes rescue line + hint
                    response_json = {
                        "response": result,  # This now includes rescue line + hint!
                        "imageData": data_url,
                        "is_end": False,
                        "is_image": True
                    }
                    print(f"[IMAGE] Response JSON keys: {list(response_json.keys())}")
                    print(f"[IMAGE] imageData field length: {len(response_json['imageData'])} chars")
                    return jsonify(response_json)
                else:
                    print(f"[IMAGE] Generation had issues: {result[:100]}...")
                    # Still return response with rescue line + hint
                    return jsonify({
                        "response": result,
                        "is_end": False,
                        "is_image": False
                    })

            # Get conversation history for this user (using IP as identifier)
            user_id = request.remote_addr
            conv_history = get_conversation_history(user_id)

            # Check for threading before generating response
            should_thread, related_turns = conv_history.check_threading(user_input)
            threaded_context = None

            if should_thread:
                threaded_context = conv_history.format_threaded_context(related_turns, user_input)
                print(f"[THREAD] Using threaded context with {len(related_turns)} previous turns")

            # Process text with threading support
            outcome = handle_text_turn(user_input, text_reply_prompt, hints_pool, rescue_lines, threaded_context)

            # Store CLEAN reply in history (without rescue lines/hints to prevent doubling)
            clean_reply = strip_rescue_and_hint(outcome)
            conv_history.add_turn(user_input, clean_reply)
            print(f"[HISTORY] Stored clean reply (length: {len(clean_reply)} chars)")

            # Check if it's a special outcome
            if outcome == "EXIT":
                return jsonify({
                    "response": "Bye! See you next time~",
                    "is_end": True
                })

            if outcome == "WIN":
                return jsonify({
                    "response": "Congratulations! You guessed the spell correctly! Happy Halloween!",
                    "is_end": True
                })

            # Check if response contains image data URL
            is_image = "data:image/" in outcome

            return jsonify({
                "response": outcome,
                "is_end": False,
                "is_image": is_image
            })

        except Exception as text_error:
            error_msg = str(text_error)

            # Check if it's a quota error
            if '402' in error_msg or 'insufficient_quota' in error_msg or 'used up your points' in error_msg:
                return jsonify({
                    "response": "[QUOTA EXCEEDED] Your POE API quota is used up. Please visit https://poe.com/api_key to add more credits. The app will work normally once credits are added!",
                    "is_end": False
                }), 200  # Return 200 to show message properly

            # Other errors
            raise text_error

    except Exception as e:
        # Safe error printing
        try:
            print(f"[ERROR] in /chat: {str(e)}")
        except:
            print("[ERROR] in /chat: [error with special characters]")

        import traceback
        traceback.print_exc()

        error_msg = str(e)

        # User-friendly error messages
        if '402' in error_msg or 'insufficient_quota' in error_msg:
            return jsonify({
                "error": "quota_exceeded",
                "response": "[QUOTA EXCEEDED] POE API credits exhausted. Visit https://poe.com/api_key to add credits.",
                "is_end": False
            }), 200
        else:
            return jsonify({
                "error": str(e),
                "response": f"[ERROR] Processing request failed: {str(e)[:100]}",
                "is_end": False
            }), 500

if __name__ == '__main__':
    print("=" * 60)
    print(" Pumpkin AI - HuggingFace Spaces")
    print("=" * 60)
    print(f"Prompt directory: {PROMPT_DIR}")
    print(f"Static files: {static_folder}")
    print(f"Hints loaded: {len(hints_pool)}")
    print(f"Rescue lines loaded: {len(rescue_lines)}")
    print("=" * 60 + "\n")

    # Get port from environment (HuggingFace Spaces sets this)
    port = int(os.environ.get('PORT', 7860))

    # Run Flask app
    app.run(
        host='0.0.0.0',
        port=port,
        debug=False  # Disable debug in production
    )