import re
import os
import time
import requests
import base64
import asyncio
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 ---
# Instead of a plain dict for conversation context, we use a dedicated class below.
user_state = {}       # e.g., { user_id: ConversationState }
conversation_context = {}  # { user_id: [ { "timestamp": ..., "role": "user"/"bot", "message": ... }, ... ] }
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"}
]

# --- Conversation State Management ---
SESSION_TIMEOUT = timedelta(minutes=5)

class ConversationState:
    def __init__(self):
        self.flow = None  # e.g., "order"
        self.step = 0
        self.data = {}
        self.last_active = datetime.utcnow()

    def update_last_active(self):
        self.last_active = datetime.utcnow()

    def is_expired(self):
        return datetime.utcnow() - self.last_active > SESSION_TIMEOUT

    def reset(self):
        self.flow = None
        self.step = 0
        self.data = {}
        self.last_active = datetime.utcnow()

# --- 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:
    """
    Implements an FSM-based order flow that:
      - In step 1, expects the user to mention a dish name (optionally with quantity)
      - In step 2, asks for quantity if not provided
      - In step 3, finalizes the order and creates a payment link
    """
    # Retrieve or initialize conversation state
    state = user_state.get(user_id)
    if state and state.is_expired():
        state.reset()
        del user_state[user_id]
        state = None

    # If the user explicitly types "order" or "menu", (re)start the order flow.
    if message.lower() in ["order", "menu"]:
        state = ConversationState()
        state.flow = "order"
        state.step = 1
        state.update_last_active()
        user_state[user_id] = state
        if message.lower() == "order":
            return "Sure! What dish would you like to order?"
        return ""

    # If no state exists but the message includes "order", start the order flow.
    if not state and "order" in message.lower():
        state = ConversationState()
        state.flow = "order"
        state.step = 1
        state.update_last_active()
        user_state[user_id] = state
        return "Sure! What dish would you like to order?"

    if state and state.flow == "order":
        state.update_last_active()
        # --- Step 1: Expecting Dish Selection (and optionally quantity) ---
        if state.step == 1:
            dish_candidates = [item["name"] for item in menu_items]
            found_dish = None
            for dish in dish_candidates:
                if dish.lower() in message.lower():
                    found_dish = dish
                    break

            numbers = re.findall(r'\d+', message)
            if found_dish:
                state.data["dish"] = found_dish
                if numbers:
                    quantity = int(numbers[0])
                    if quantity <= 0:
                        return "Please enter a valid quantity (e.g., 1, 2, 3)."
                    state.data["quantity"] = quantity
                    state.step = 3  # ready to finalize the order
                else:
                    state.step = 2  # ask for quantity
            else:
                return "I couldn't identify the dish. Please type the dish name from our menu."

        # --- Step 2: Asking for Quantity ---
        if state.step == 2:
            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)."
            state.data["quantity"] = quantity
            state.step = 3

        # --- Step 3: Finalize Order ---
        if state.step == 3:
            order_id = f"ORD-{int(time.time())}"
            state.data["order_id"] = order_id
            price_per_serving = 1500  # fixed price per serving for demonstration
            quantity = state.data.get("quantity", 1)
            total_price = quantity * price_per_serving
            state.data["price"] = str(total_price)

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

            email = "customer@example.com"  # Placeholder; retrieve from profile if available
            payment_data = create_paystack_payment_link(email, total_price * 100, order_id)
            dish_name = state.data.get("dish", "")
            state.reset()
            if user_id in user_state:
                del user_state[user_id]

            if payment_data.get("status"):
                payment_link = payment_data["data"]["authorization_url"]
                return (f"Thank you for your order of {quantity} serving(s) of {dish_name}! "
                        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."

    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.")

    # Initialize conversation context for the user if not present.
    if user_id not in conversation_context:
        conversation_context[user_id] = []
    # Append the inbound message to the conversation context.
    conversation_context[user_id].append({
        "timestamp": datetime.utcnow().isoformat(),
        "role": "user",
        "message": user_message
    })

    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)
        conversation_context[user_id].append({
            "timestamp": datetime.utcnow().isoformat(),
            "role": "bot",
            "message": order_response
        })
        return JSONResponse(content={"response": sentiment_modifier + order_response})

    # --- Menu Display ---
    if "menu" in user_message.lower():
        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))
        conversation_context[user_id].append({
            "timestamp": datetime.utcnow().isoformat(),
            "role": "bot",
            "message": response_payload["response"]
        })
        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:
            state = ConversationState()
            state.flow = "order"
            # **** FIX: Set step to 2 since the dish is already selected ****
            state.step = 2  
            state.data["dish"] = selected_dish
            state.update_last_active()
            user_state[user_id] = state
            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)
            conversation_context[user_id].append({
                "timestamp": datetime.utcnow().isoformat(),
                "role": "bot",
                "message": 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)
            conversation_context[user_id].append({
                "timestamp": datetime.utcnow().isoformat(),
                "role": "bot",
                "message": 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)
        conversation_context[user_id].append({
            "timestamp": datetime.utcnow().isoformat(),
            "role": "bot",
            "message": response_text
        })
        return JSONResponse(content={"response": sentiment_modifier + response_text})

    # --- Fallback: LLM Response Streaming with Conversation Context ---
    recent_context = conversation_context.get(user_id, [])[-5:]
    context_str = "\n".join([f"{entry['role'].capitalize()}: {entry['message']}" for entry in recent_context])
    prompt = f"Conversation context:\n{context_str}\nUser query: {user_message}\nGenerate a helpful, personalized response for a restaurant chatbot."
    def stream_response():
        for chunk in stream_text_completion(prompt):
            yield chunk
    fallback_log = f"LLM fallback response for prompt: {prompt}"
    background_tasks.add_task(log_chat_to_db, user_id, "outbound", fallback_log)
    return StreamingResponse(stream_response(), media_type="text/plain")

# --- Other Endpoints (Chat History, Order Details, User Profile, Analytics, Voice, Payment Callback) ---
# ... (Implement other endpoints as needed) ...

@app.get("/chat_history/{user_id}")
async def get_chat_history(user_id: str):
    async with async_session() as session:
        result = await session.execute(
            ChatHistory.__table__.select().where(ChatHistory.user_id == user_id)
        )
        history = result.fetchall()
        return [dict(row) for row in history]

@app.get("/order/{order_id}")
async def get_order(order_id: str):
    async with async_session() as session:
        result = await session.execute(
            Order.__table__.select().where(Order.order_id == order_id)
        )
        order = result.fetchone()
        if order:
            return dict(order)
        else:
            raise HTTPException(status_code=404, detail="Order not found.")

@app.get("/user_profile/{user_id}")
async def get_user_profile(user_id: str):
    profile = await get_or_create_user_profile(user_id)
    return {
        "user_id": profile.user_id,
        "phone_number": profile.phone_number,
        "name": profile.name,
        "email": profile.email,
        "preferences": profile.preferences,
        "last_interaction": profile.last_interaction.isoformat()
    }

@app.get("/analytics")
async def get_analytics():
    async with async_session() as session:
        msg_result = await session.execute(ChatHistory.__table__.count())
        total_messages = msg_result.scalar() or 0
        order_result = await session.execute(Order.__table__.count())
        total_orders = order_result.scalar() or 0
        sentiment_result = await session.execute("SELECT AVG(sentiment_score) FROM sentiment_logs")
        avg_sentiment = sentiment_result.scalar() or 0
    return {
        "total_messages": total_messages,
        "total_orders": total_orders,
        "average_sentiment": avg_sentiment
    }

@app.post("/voice")
async def process_voice(file: UploadFile = File(...)):
    contents = await file.read()
    simulated_text = "Simulated speech-to-text conversion result."
    return {"transcription": simulated_text}

@app.post("/payment_callback")
async def payment_callback(request: Request):
    data = await request.json()
    order_id = data.get("reference")
    new_status = data.get("status", "Paid")
    async with async_session() as session:
        result = await session.execute(
            Order.__table__.select().where(Order.order_id == order_id)
        )
        order = result.scalar_one_or_none()
        if order:
            order.status = new_status
            await session.commit()
            return JSONResponse(content={"message": "Order updated successfully."})
        else:
            raise HTTPException(status_code=404, detail="Order not found.")

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