import re
import os
import time
import requests
import base64
from datetime import datetime, timedelta
from bs4 import BeautifulSoup
from sqlalchemy import select

from fastapi import FastAPI, Request, HTTPException, BackgroundTasks, UploadFile, File, Form
from fastapi.responses import JSONResponse, StreamingResponse

import openai

# For sentiment analysis using TextBlob
from textblob import TextBlob

# SQLAlchemy Imports (Async)
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy import Column, Integer, String, DateTime, Text, Float

# --- Configuration & Environment Variables ---
SPOONACULAR_API_KEY = os.getenv("SPOONACULAR_API_KEY", "815bf76e0764456293f0e96e080e8f60")
PAYSTACK_SECRET_KEY = os.getenv("PAYSTACK_SECRET_KEY", "sk_test_52a354dba436437c3ea86c9089c640ad12a7b115")
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql+asyncpg://postgres.lgbnxplydqdymepehirg:Lovyelias5584.@aws-0-eu-central-1.pooler.supabase.com:5432/postgres")
NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY", "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")
openai.api_key = os.getenv("OPENAI_API_KEY", "your_openai_api_key")

# --- Database Setup ---
Base = declarative_base()

class ChatHistory(Base):
    __tablename__ = "chat_history"
    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(String, index=True)
    timestamp = Column(DateTime, default=datetime.utcnow)
    direction = Column(String)  # 'inbound' or 'outbound'
    message = Column(Text)

class Order(Base):
    __tablename__ = "orders"
    id = Column(Integer, primary_key=True, index=True)
    order_id = Column(String, unique=True, index=True)
    user_id = Column(String, index=True)
    dish = Column(String)
    quantity = Column(String)
    price = Column(String, default="0")
    status = Column(String, default="Pending Payment")
    payment_reference = Column(String, nullable=True)
    timestamp = Column(DateTime, default=datetime.utcnow)

class UserProfile(Base):
    __tablename__ = "user_profiles"
    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(String, unique=True, index=True)
    phone_number = Column(String, unique=True, index=True, nullable=True)
    name = Column(String, default="Valued Customer")
    email = Column(String, default="unknown@example.com")
    preferences = Column(Text, default="")
    last_interaction = Column(DateTime, default=datetime.utcnow)

class SentimentLog(Base):
    __tablename__ = "sentiment_logs"
    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(String, index=True)
    timestamp = Column(DateTime, default=datetime.utcnow)
    sentiment_score = Column(Float)
    message = Column(Text)

engine = create_async_engine(DATABASE_URL, echo=True)
async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

async def init_db():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

# --- Global In-Memory Stores ---
user_state = {}       # e.g., { user_id: { "flow": str, "step": int, "data": dict, "last_active": datetime } }
conversation_context = {}  
proactive_timer = {}   

menu_items = [
    {"name": "Jollof Rice", "description": "A spicy and flavorful rice dish", "price": 1500, "nutrition": "Calories: 300 kcal, Carbs: 50g, Protein: 10g, Fat: 5g"},
    {"name": "Fried Rice", "description": "A savory rice dish with vegetables and meat", "price": 1200, "nutrition": "Calories: 350 kcal, Carbs: 55g, Protein: 12g, Fat: 8g"},
    {"name": "Chicken Wings", "description": "Crispy fried chicken wings", "price": 2000, "nutrition": "Calories: 400 kcal, Carbs: 20g, Protein: 25g, Fat: 15g"},
    {"name": "Egusi Soup", "description": "A rich and hearty soup made with melon seeds", "price": 1000, "nutrition": "Calories: 250 kcal, Carbs: 15g, Protein: 8g, Fat: 10g"}
]

# --- Utility Functions ---
async def log_chat_to_db(user_id: str, direction: str, message: str):
    async with async_session() as session:
        entry = ChatHistory(user_id=user_id, direction=direction, message=message)
        session.add(entry)
        await session.commit()

async def log_sentiment(user_id: str, message: str, score: float):
    async with async_session() as session:
        entry = SentimentLog(user_id=user_id, sentiment_score=score, message=message)
        session.add(entry)
        await session.commit()

def analyze_sentiment(text: str) -> float:
    blob = TextBlob(text)
    return blob.sentiment.polarity

def google_image_scrape(query: str) -> str:
    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"}
    search_url = f"https://www.google.com/search?tbm=isch&q={query}"
    try:
        response = requests.get(search_url, headers=headers, timeout=5)
    except Exception:
        return ""
    if response.status_code == 200:
        soup = BeautifulSoup(response.text, "html.parser")
        img_tags = soup.find_all("img")
        for img in img_tags:
            src = img.get("src")
            if src and src.startswith("http"):
                return src
    return ""

def create_paystack_payment_link(email: str, amount: int, reference: str) -> dict:
    url = "https://api.paystack.co/transaction/initialize"
    headers = {
        "Authorization": f"Bearer {PAYSTACK_SECRET_KEY}",
        "Content-Type": "application/json",
    }
    data = {
        "email": email,
        "amount": amount,
        "reference": reference,
        "callback_url": "https://yourdomain.com/payment_callback"
    }
    try:
        response = requests.post(url, json=data, headers=headers, timeout=10)
        if response.status_code == 200:
            return response.json()
        else:
            return {"status": False, "message": "Failed to initialize payment."}
    except Exception as e:
        return {"status": False, "message": str(e)}

# --- NVIDIA LLM Streaming Functions ---
def stream_text_completion(prompt: str):
    from openai import OpenAI
    client = OpenAI(
        base_url="https://integrate.api.nvidia.com/v1",
        api_key=NVIDIA_API_KEY
    )
    completion = client.chat.completions.create(
        model="meta/llama-3.1-405b-instruct",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.2,
        top_p=0.7,
        max_tokens=1024,
        stream=True
    )
    for chunk in completion:
        if chunk.choices[0].delta.content is not None:
            yield chunk.choices[0].delta.content

def stream_image_completion(image_b64: str):
    invoke_url = "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct/chat/completions"
    headers = {
        "Authorization": f"Bearer {NVIDIA_API_KEY}",
        "Accept": "text/event-stream"
    }
    payload = {
        "model": "meta/llama-3.2-90b-vision-instruct",
        "messages": [
            {
                "role": "user",
                "content": f'What is in this image? <img src="data:image/png;base64,{image_b64}" />'
            }
        ],
        "max_tokens": 512,
        "temperature": 1.00,
        "top_p": 1.00,
        "stream": True
    }
    response = requests.post(invoke_url, headers=headers, json=payload, stream=True)
    for line in response.iter_lines():
        if line:
            yield line.decode("utf-8") + "\n"

# --- Advanced Internal Flow: Order Processing & Payment Integration ---
def process_order_flow(user_id: str, message: str) -> str:
    # If user types a command that should override any active order flow:
    if message.lower() in ["order", "menu"]:
        if user_id in user_state:
            del user_state[user_id]
        if message.lower() == "order":
            user_state[user_id] = {"flow": "order", "step": 1, "data": {}, "last_active": datetime.utcnow()}
            return "Sure! What dish would you like to order?"
        # For "menu", we return empty to allow the menu-handling code to run.
        return ""

    if user_id in user_state:
        state = user_state[user_id]
        flow = state.get("flow")
        step = state.get("step")
        data = state.get("data", {})
        if flow == "order":
            if step == 1:
                # Validate dish name (exact match is used)
                dish_name = message.title()
                dish = next((item for item in menu_items if item["name"].lower() == dish_name.lower()), None)
                if dish:
                    data["dish"] = dish_name
                    state["step"] = 2
                    return f"You selected {data['dish']}. How many servings would you like?"
                else:
                    return f"Sorry, we don't have {dish_name} on the menu. Please choose another dish."
            elif step == 2:
                # Extract the first number from the message (e.g., '1 servings' becomes 1)
                numbers = re.findall(r'\d+', message)
                if not numbers:
                    return "Please enter a valid number for the quantity (e.g., 1, 2, 3)."
                quantity = int(numbers[0])
                if quantity <= 0:
                    return "Please enter a valid quantity (e.g., 1, 2, 3)."
                data["quantity"] = quantity
                order_id = f"ORD-{int(time.time())}"
                data["order_id"] = order_id
                price_per_serving = 1500  # ₦1500 per serving
                total_price = quantity * price_per_serving
                data["price"] = str(total_price)

                # Save order asynchronously
                import asyncio
                async def save_order():
                    async with async_session() as session:
                        order = Order(
                            order_id=order_id,
                            user_id=user_id,
                            dish=data["dish"],
                            quantity=str(quantity),
                            price=str(total_price),
                            status="Pending Payment"
                        )
                        session.add(order)
                        await session.commit()
                asyncio.create_task(save_order())

                # Clear conversation state for order flow.
                del user_state[user_id]

                # Retrieve email from user profile if available (using a placeholder here)
                email = "customer@example.com"
                payment_data = create_paystack_payment_link(email, total_price * 100, order_id)
                if payment_data.get("status"):
                    payment_link = payment_data["data"]["authorization_url"]
                    return (f"Thank you for your order of {data['quantity']} serving(s) of {data['dish']}! "
                            f"Your Order ID is {order_id}.\nPlease complete payment here: {payment_link}")
                else:
                    return f"Your order has been placed with Order ID {order_id}, but we could not initialize payment. Please try again later."
    else:
        if "order" in message.lower():
            user_state[user_id] = {"flow": "order", "step": 1, "data": {}, "last_active": datetime.utcnow()}
            return "Sure! What dish would you like to order?"
    return ""

# --- User Profile Functions ---
async def get_or_create_user_profile(user_id: str, phone_number: str = None) -> UserProfile:
    async with async_session() as session:
        result = await session.execute(
            select(UserProfile).where(UserProfile.user_id == user_id)
        )
        profile = result.scalars().first()
        if profile is None:
            profile = UserProfile(
                user_id=user_id,
                phone_number=phone_number,
                last_interaction=datetime.utcnow()
            )
            session.add(profile)
            await session.commit()
        return profile

async def update_user_last_interaction(user_id: str):
    async with async_session() as session:
        result = await session.execute(
            select(UserProfile).where(UserProfile.user_id == user_id)
        )
        profile = result.scalars().first()
        if profile:
            profile.last_interaction = datetime.utcnow()
            await session.commit()

# --- Proactive Engagement: Warm Greetings ---
async def send_proactive_greeting(user_id: str):
    greeting = "Hi again! We miss you. Would you like to see our new menu items or get personalized recommendations?"
    await log_chat_to_db(user_id, "outbound", greeting)
    return greeting

# --- FastAPI Setup & Endpoints ---
app = FastAPI()

@app.on_event("startup")
async def on_startup():
    await init_db()

@app.post("/chatbot")
async def chatbot_response(request: Request, background_tasks: BackgroundTasks):
    data = await request.json()
    user_id = data.get("user_id")
    phone_number = data.get("phone_number")
    user_message = data.get("message", "").strip()
    is_image = data.get("is_image", False)
    image_b64 = data.get("image_base64", None)

    if not user_id:
        raise HTTPException(status_code=400, detail="Missing user_id in payload.")

    background_tasks.add_task(log_chat_to_db, user_id, "inbound", user_message)
    await update_user_last_interaction(user_id)
    await get_or_create_user_profile(user_id, phone_number)

    # Handle image queries
    if is_image and image_b64:
        if len(image_b64) >= 180_000:
            raise HTTPException(status_code=400, detail="Image too large.")
        return StreamingResponse(stream_image_completion(image_b64), media_type="text/plain")

    sentiment_score = analyze_sentiment(user_message)
    background_tasks.add_task(log_sentiment, user_id, user_message, sentiment_score)
    sentiment_modifier = ""
    if sentiment_score < -0.3:
        sentiment_modifier = "I'm sorry if you're having a tough time. "
    elif sentiment_score > 0.3:
        sentiment_modifier = "Great to hear from you! "

    # --- Order Flow Handling ---
    order_response = process_order_flow(user_id, user_message)
    if order_response:
        background_tasks.add_task(log_chat_to_db, user_id, "outbound", order_response)
        return JSONResponse(content={"response": sentiment_modifier + order_response})

    # --- Menu Display ---
    if "menu" in user_message.lower():
        # Clear any active order flow if user explicitly asks for menu.
        if user_id in user_state:
            del user_state[user_id]
        menu_with_images = []
        for index, item in enumerate(menu_items, start=1):
            image_url = google_image_scrape(item["name"])
            menu_with_images.append({
                "number": index,
                "name": item["name"],
                "description": item["description"],
                "price": item["price"],
                "image_url": image_url
            })
        response_payload = {
            "response": sentiment_modifier + "Here’s our delicious menu:",
            "menu": menu_with_images,
            "follow_up": (
                "To order, type the *number* or *name* of the dish you'd like. "
                "For example, type '1' or 'Jollof Rice' to order Jollof Rice.\n\n"
                "You can also ask for nutritional facts by typing, for example, 'Nutritional facts for Jollof Rice'."
            )
        }
        background_tasks.add_task(log_chat_to_db, user_id, "outbound", str(response_payload))
        return JSONResponse(content=response_payload)

    # --- Dish Selection via Menu ---
    if any(item["name"].lower() in user_message.lower() for item in menu_items) or \
       any(str(index) == user_message.strip() for index, item in enumerate(menu_items, start=1)):
        selected_dish = None
        if user_message.strip().isdigit():
            dish_number = int(user_message.strip())
            if 1 <= dish_number <= len(menu_items):
                selected_dish = menu_items[dish_number - 1]["name"]
        else:
            for item in menu_items:
                if item["name"].lower() in user_message.lower():
                    selected_dish = item["name"]
                    break
        if selected_dish:
            # Trigger a new order flow
            user_state[user_id] = {"flow": "order", "step": 1, "data": {"dish": selected_dish}, "last_active": datetime.utcnow()}
            response_text = f"You selected {selected_dish}. How many servings would you like?"
            background_tasks.add_task(log_chat_to_db, user_id, "outbound", response_text)
            return JSONResponse(content={"response": sentiment_modifier + response_text})
        else:
            response_text = "Sorry, I couldn't find that dish in the menu. Please try again."
            background_tasks.add_task(log_chat_to_db, user_id, "outbound", response_text)
            return JSONResponse(content={"response": sentiment_modifier + response_text})

    # --- Nutritional Facts ---
    if "nutritional facts for" in user_message.lower():
        dish_name = user_message.lower().replace("nutritional facts for", "").strip().title()
        dish = next((item for item in menu_items if item["name"].lower() == dish_name.lower()), None)
        if dish:
            response_text = f"Nutritional facts for {dish['name']}:\n{dish['nutrition']}"
        else:
            response_text = f"Sorry, I couldn't find nutritional facts for {dish_name}."
        background_tasks.add_task(log_chat_to_db, user_id, "outbound", response_text)
        return JSONResponse(content={"response": sentiment_modifier + response_text})

    # --- Fallback: LLM Response Streaming ---
    prompt = f"User query: {user_message}\nGenerate a helpful, personalized response for a restaurant chatbot."
    def stream_response():
        for chunk in stream_text_completion(prompt):
            yield chunk
    background_tasks.add_task(log_chat_to_db, user_id, "outbound", f"LLM fallback response for prompt: {prompt}")
    return StreamingResponse(stream_response(), media_type="text/plain")
    
# --- Other Endpoints (Chat History, Order Details, User Profile, Analytics, Voice, Payment Callback) ---
# ... (unchanged for brevity) ...

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)