api.py

from fastapi import FastAPI, HTTPException, Depends, status
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from fastapi.middleware.cors import CORSMiddleware
from sqlalchemy.orm import Session, joinedload
from sqlalchemy import func
from datetime import datetime, timedelta
from typing import Optional, List
from jose import JWTError, jwt
from passlib.context import CryptContext
from pydantic import BaseModel
import uvicorn
import os
import base64

from chat_engine import ChatEngine
from image_engine import ImageEngine
import models
import schemas
import firebase_admin
from firebase_admin import credentials, firestore, auth

# Initialize Firebase Admin
if not firebase_admin._apps:
    if os.path.exists("serviceAccountKey.json"):
        cred = credentials.Certificate("serviceAccountKey.json")
    else:
        # Try getting from env var (for Hugging Face)
        key_json = os.environ.get("FIREBASE_SERVICE_ACCOUNT_KEY")
        if key_json:
            import json
            cred_dict = json.loads(key_json)
            cred = credentials.Certificate(cred_dict)
        else:
            print("Warning: No service account key found. Firebase features will fail.")
            cred = None
            
    if cred:
        firebase_admin.initialize_app(cred)

if firebase_admin._apps:
    db = firestore.client()
else:
    db = None

# Initialize engines
print("Initializing AI Engines...")
chat_engine = ChatEngine()
image_engine = ImageEngine()
rag_engine = RAGEngine()
print("AI Engines Ready!")

# Auth Dependency
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

async def get_current_user(token: str = Depends(oauth2_scheme)):
    try:
        decoded_token = auth.verify_id_token(token)
        uid = decoded_token['uid']
        # Get user data from Firestore
        user_doc = db.collection('users').document(uid).get()
        if not user_doc.exists:
            # Create user if not exists (first login)
            user_data = {
                "email": decoded_token.get('email'),
                "full_name": decoded_token.get('name', 'User'),
                "created_at": datetime.utcnow(),
                "is_admin": False
            }
            db.collection('users').document(uid).set(user_data)
            return {**user_data, "id": uid}
        
        return {**user_doc.to_dict(), "id": uid}
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=f"Invalid authentication credentials: {str(e)}",
            headers={"WWW-Authenticate": "Bearer"},
        )

async def get_current_admin(current_user: dict = Depends(get_current_user)):
    if not current_user.get("is_admin"):
        raise HTTPException(status_code=403, detail="Not authorized")
    return current_user

# Auth Endpoints
# Note: Registration and Login are handled by Firebase on the Frontend.
# The backend only verifies the ID token via get_current_user.

@app.get("/users/me")
async def read_users_me(current_user: dict = Depends(get_current_user)):
    return current_user

# Conversation Endpoints
@app.post("/conversations")
async def create_conversation(conversation: schemas.ConversationCreate, current_user: dict = Depends(get_current_user)):
    try:
        new_conv_ref = db.collection('conversations').document()
        conv_data = {
            "id": new_conv_ref.id,
            "user_id": current_user['id'],
            "title": conversation.title,
            "created_at": datetime.utcnow(),
            "updated_at": datetime.utcnow()
        }
        new_conv_ref.set(conv_data)
        return conv_data
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/conversations")
async def get_conversations(current_user: dict = Depends(get_current_user)):
    try:
        docs = db.collection('conversations').where('user_id', '==', current_user['id']).order_by('updated_at', direction=firestore.Query.DESCENDING).stream()
        return [doc.to_dict() for doc in docs]
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/conversations/{conversation_id}/messages")
async def get_conversation_messages(conversation_id: str, current_user: dict = Depends(get_current_user)):
    try:
        # Verify ownership
        conv_ref = db.collection('conversations').document(conversation_id)
        conv = conv_ref.get()
        if not conv.exists or conv.to_dict()['user_id'] != current_user['id']:
            raise HTTPException(status_code=404, detail="Conversation not found")
            
        msgs = conv_ref.collection('messages').order_by('timestamp').stream()
        return [msg.to_dict() for msg in msgs]
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Saved Prompt Endpoints
@app.post("/prompts")
async def create_prompt(prompt: schemas.SavedPromptCreate, current_user: dict = Depends(get_current_user)):
    try:
        new_prompt_ref = db.collection('prompts').document()
        prompt_data = {
            "id": new_prompt_ref.id,
            "user_id": current_user['id'],
            "title": prompt.title,
            "content": prompt.content,
            "tags": prompt.tags,
            "created_at": datetime.utcnow()
        }
        new_prompt_ref.set(prompt_data)
        return prompt_data
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/prompts")
async def get_prompts(current_user: dict = Depends(get_current_user)):
    try:
        docs = db.collection('prompts').where('user_id', '==', current_user['id']).order_by('created_at', direction=firestore.Query.DESCENDING).stream()
        return [doc.to_dict() for doc in docs]
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/prompts/{prompt_id}")
async def delete_prompt(prompt_id: str, current_user: dict = Depends(get_current_user)):
    try:
        prompt_ref = db.collection('prompts').document(prompt_id)
        prompt = prompt_ref.get()
        if not prompt.exists or prompt.to_dict()['user_id'] != current_user['id']:
            raise HTTPException(status_code=404, detail="Prompt not found")
        prompt_ref.delete()
        return {"status": "success"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Admin Endpoints
@app.get("/admin/users")
async def get_all_users(current_user: dict = Depends(get_current_admin)):
    try:
        users = db.collection('users').stream()
        result = []
        for user in users:
            user_data = user.to_dict()
            # Count messages (this might be expensive in Firestore, maybe skip or approximate)
            # For now, let's just return user data
            result.append(user_data)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/admin/activity")
async def get_all_activity(current_user: dict = Depends(get_current_admin)):
    # This is hard in Firestore without a global collection group query
    # For now, return empty or implement a specific 'activity' log collection
    return []

# Protected AI Endpoints
class ChatRequest(BaseModel):
    message: str
    history: list = []
    language: str = "English"
    conversation_id: Optional[int] = None

class ImageRequest(BaseModel):
    prompt: str

@app.get("/")
def read_root():
    return {"status": "Backend is running", "message": "Go to /docs to see the API"}

@app.post("/chat")
async def chat(request: ChatRequest, current_user: dict = Depends(get_current_user)):
    # ... (Keep existing /chat for backward compatibility if needed, or redirect logic)
    # For now, let's keep /chat as blocking and add /chat/stream
    try:
        # Generate Response
        response = chat_engine.generate_response(request.message, request.history)
        
        # Save to Firestore if conversation_id is present
        if request.conversation_id:
            conv_ref = db.collection('conversations').document(request.conversation_id)
            # User Msg
            conv_ref.collection('messages').add({
                "role": "user",
                "content": request.message,
                "timestamp": datetime.utcnow()
            })
            # AI Msg
            conv_ref.collection('messages').add({
                "role": "assistant",
                "content": response,
                "timestamp": datetime.utcnow()
            })
            conv_ref.update({"updated_at": datetime.utcnow()})
        
        return {"response": response}
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

# RAG Endpoints
@app.post("/upload")
async def upload_file(file: UploadFile = File(...), current_user: dict = Depends(get_current_user)):
    try:
        # Save file locally
        upload_dir = "uploads"
        os.makedirs(upload_dir, exist_ok=True)
        file_path = os.path.join(upload_dir, file.filename)
        
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
            
        # Ingest into RAG
        rag_engine.ingest_file(file_path)
        
        return {"filename": file.filename, "status": "ingested"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/chat/stream")
async def chat_stream(request: ChatRequest, current_user: dict = Depends(get_current_user)):
    try:
        # Check for RAG context
        context = ""
        rag_docs = rag_engine.search(request.message)
        if rag_docs:
            context = "\n\nRelevant Context:\n" + "\n".join(rag_docs) + "\n\n"
            print(f"Found {len(rag_docs)} relevant documents.")

        # Save User Message
        if request.conversation_id:
            conv_ref = db.collection('conversations').document(request.conversation_id)
            conv_ref.collection('messages').add({
                "role": "user",
                "content": request.message,
                "timestamp": datetime.utcnow()
            })
            conv_ref.update({"updated_at": datetime.utcnow()})

        async def stream_generator():
            full_response = ""
            # Prepend context to the message sent to AI (but not saved in DB as user message)
            augmented_message = context + request.message if context else request.message
            
            for token in chat_engine.generate_stream(augmented_message, request.history, request.language):
                full_response += token
                yield token
            
            # Save AI Message after generation
            if request.conversation_id:
                conv_ref = db.collection('conversations').document(request.conversation_id)
                conv_ref.collection('messages').add({
                    "role": "assistant",
                    "content": full_response,
                    "timestamp": datetime.utcnow()
                })
            
            print(f"Generated response for conv {request.conversation_id}")

        return StreamingResponse(stream_generator(), media_type="text/plain")

    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/generate-image")
async def generate_image(request: ImageRequest, current_user: dict = Depends(get_current_user)):
    try:
        # Generate image to a temporary file
        filename = "temp_generated.png"
        image_engine.generate_image(request.prompt, output_path=filename)
        
        # Read and encode to base64 to send to frontend
        with open(filename, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
            
        return {"image_base64": encoded_string}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

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