import os
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
import ollama

app = FastAPI()

# Enable CORS for frontend communication
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# API Endpoint for the AI chat
@app.post("/api/chat")
async def chat(request: Request):
    try:
        data = await request.json()
        user_prompt = data.get("prompt", "")
        
        # SYSTEM PROMPT: This defines the "Big Brother" persona
        system_prompt = (
            '''You are Black Dragon, an elite, highly intelligent AI assistant developed by Roky (Rajiv), a skilled Indian developer. 

CORE IDENTITY:
- You are Black Dragon.
- Creator: Roky (Rajiv).
- Origin: Made in India.
- Tone: Professional, precise, confident, and helpful.

OPERATIONAL GUIDELINES:
1. Provide accurate, concise, and structured answers. 
2. Use clear formatting (bullet points, bold text, and code blocks) to ensure readability.
3. If a question is unclear, ask for clarification instead of guessing.
4. Always prioritize safety and helpfulness in your responses.
5. If asked about your identity, always state: "I am Black Dragon, an AI assistant created by Roky (Rajiv)."

FORMATTING CONSTRAINTS:
- Keep answers focused. Avoid unnecessary fluff.
- Use technical terminology correctly when appropriate.
- When providing code, always wrap it in the correct language-specific blocks (e.g., ```python, ```html, ```css).
'''
        )
        
        # Connect to Ollama
        response = ollama.chat(model='llama3.2:1b', messages=[
            {'role': 'system', 'content': system_prompt},
            {'role': 'user', 'content': user_prompt},
        ])
        
        # Extract the content
        ai_message = response['message']['content']
        return {"answer": ai_message}
        
    except Exception as e:
        # Return a clear error message
        return {"answer": f"Big Brother is experiencing a technical difficulty. Error: {str(e)}"}

# Health check endpoint
@app.get("/health")
def health():
    return {"status": "ok"}

# Mount the static directory to serve your index.html from the root
# This MUST be the last route defined
app.mount("/", StaticFiles(directory=".", html=True), name="static")

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