main.py

import uuid
import json
import asyncio
import httpx
import sqlite3
import time
from pathlib import Path
from datetime import datetime
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, UploadFile, File, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.middleware.cors import CORSMiddleware
from typing import List, Dict, Set

app = FastAPI(title="Antaram Chat AI Pro")

# --- Configuration ---
AI_URL = "https://sarveshpatel-unsloth0-6bqwen.hf.space/chat/raw"
MAX_TOKENS = 30000
SYSTEM_PROMPT = "You are Antaram AI. Answer concisely, accurately, and professionally. Use emojis very sparingly (max 1 per response). Do not hallucinate."

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Setup
BASE_DIR = Path(__file__).resolve().parent
UPLOAD_DIR = BASE_DIR / "uploads"
UPLOAD_DIR.mkdir(exist_ok=True)

templates = Jinja2Templates(directory="templates")
app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")

# --- Database Setup (SQLite) ---
DB_PATH = "chat.db"

def init_db():
    with sqlite3.connect(DB_PATH) as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS messages (
                id TEXT PRIMARY KEY,
                room_id TEXT,
                username TEXT,
                content TEXT,
                file_url TEXT,
                file_type TEXT,
                reply_to_id TEXT,
                reply_content TEXT,
                timestamp REAL
            )
        """)
        conn.commit()

init_db()

# --- State ---
active_rooms: Dict[str, List[WebSocket]] = {}
active_users: Dict[str, Set[str]] = {} # room_id -> set of usernames

# --- Helpers ---

def save_message(room_id, username, content, file_data=None, reply_to_id=None, reply_content=None):
    msg_id = str(uuid.uuid4())
    ts = time.time()
    file_url = file_data['file_url'] if file_data else None
    file_type = file_data['file_type'] if file_data else None
    
    with sqlite3.connect(DB_PATH) as conn:
        conn.execute(
            "INSERT INTO messages VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (msg_id, room_id, username, content, file_url, file_type, reply_to_id, reply_content, ts)
        )
    return {
        "id": msg_id, "username": username, "text": content, 
        "file": file_data, "reply_to": reply_to_id, "reply_content": reply_content,
        "timestamp": ts, "type": "message"
    }

def get_history(room_id, limit=50):
    with sqlite3.connect(DB_PATH) as conn:
        conn.row_factory = sqlite3.Row
        cursor = conn.execute(
            "SELECT * FROM messages WHERE room_id = ? ORDER BY timestamp ASC LIMIT ?", 
            (room_id, limit)
        )
        rows = cursor.fetchall()
        
    history = []
    for row in rows:
        file_data = None
        if row['file_url']:
            file_data = {"file_url": row['file_url'], "file_type": row['file_type'], "original_name": "File"}
            
        history.append({
            "type": "message",
            "id": row['id'],
            "username": row['username'],
            "text": row['content'],
            "file": file_data,
            "reply_to": row['reply_to_id'],
            "reply_content": row['reply_content'],
            "timestamp": row['timestamp']
        })
    return history

# --- AI Logic ---

async def stream_antaram_ai(room_id: str, prompt: str, context_msg: str = None):
    """Streams AI response. If context_msg is provided (reply), it's added to prompt."""
    
    # 1. Notify Start
    await broadcast_to_room(room_id, json.dumps({"type": "ai_start", "username": "Antaram AI"}))

    final_prompt = f"{SYSTEM_PROMPT}\n\n"
    if context_msg:
        final_prompt += f"Context (User replied to this): {context_msg}\n"
    
    clean_user_prompt = prompt.replace("@antaram.ai", "").strip()
    final_prompt += f"User Query: {clean_user_prompt}"

    full_response = ""
    
    try:
        async with httpx.AsyncClient(timeout=60) as client:
            async with client.stream("POST", AI_URL, json={"prompt": final_prompt, "max_tokens": MAX_TOKENS}) as response:
                async for chunk in response.aiter_text():
                    clean_chunk = chunk.replace("<think>", "").replace("</think>", "")
                    if clean_chunk:
                        full_response += clean_chunk
                        await broadcast_to_room(room_id, json.dumps({
                            "type": "ai_chunk", "chunk": clean_chunk
                        }))
    except Exception as e:
        await broadcast_to_room(room_id, json.dumps({"type": "error", "message": "AI Unreachable"}))

    # 2. Notify End & Save AI response to DB
    save_message(room_id, "Antaram AI", full_response)
    await broadcast_to_room(room_id, json.dumps({"type": "ai_end"}))

# --- Routes ---

@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    return templates.TemplateResponse("index.html", {"request": request, "room_id": None})

@app.get("/room/{room_id}", response_class=HTMLResponse)
async def dynamic_room(request: Request, room_id: str):
    return templates.TemplateResponse("index.html", {"request": request, "room_id": room_id.upper()})

@app.post("/create-room")
async def create_room():
    room_id = str(uuid.uuid4())[:8].upper()
    return {"room_id": room_id, "success": True}

@app.post("/upload-file/{room_id}")
async def upload_file(room_id: str, file: UploadFile = File(...)):
    file_ext = Path(file.filename).suffix
    unique_name = f"{uuid.uuid4().hex}{file_ext}"
    file_path = UPLOAD_DIR / unique_name
    content = await file.read()
    with open(file_path, "wb") as f:
        f.write(content)
    
    return {"success": True, "file_info": {
        "file_url": f"/uploads/{unique_name}",
        "file_type": file.content_type,
        "original_name": file.filename
    }}

# --- WebSocket ---

@app.websocket("/ws/{room_id}")
async def websocket_endpoint(websocket: WebSocket, room_id: str):
    room_id = room_id.upper()
    await websocket.accept()
    
    if room_id not in active_rooms:
        active_rooms[room_id] = []
        active_users[room_id] = set()

    active_rooms[room_id].append(websocket)
    
    # 1. Send History First
    history = get_history(room_id)
    await websocket.send_text(json.dumps({"type": "history", "data": history}))
    
    try:
        while True:
            data = await websocket.receive_text()
            msg_data = json.loads(data)
            
            msg_type = msg_data.get("type")
            username = msg_data.get("username", "Guest")
            
            # Handle Join (to update user list)
            if msg_type == "join":
                active_users[room_id].add(username)
                await broadcast_to_room(room_id, json.dumps({
                    "type": "system", 
                    "message": f"{username} joined.",
                    "users": list(active_users[room_id])
                }))
                continue

            # Handle Message
            if msg_type == "message":
                text = msg_data.get("text", "")
                reply_to = msg_data.get("reply_to")
                reply_content = msg_data.get("reply_content")
                
                # Save to DB
                saved_msg = save_message(room_id, username, text, msg_data.get("file"), reply_to, reply_content)
                
                # Broadcast
                await broadcast_to_room(room_id, json.dumps(saved_msg))

                # Check AI
                if "@antaram.ai" in text.lower():
                    # If user is replying to a message AND mentioning AI, pass that context
                    asyncio.create_task(stream_antaram_ai(room_id, text, reply_content))

    except WebSocketDisconnect:
        if room_id in active_rooms:
            if websocket in active_rooms[room_id]:
                active_rooms[room_id].remove(websocket)
            # We don't remove user immediately from set to keep history context, 
            # or you can remove if you want strict online status.
            await broadcast_to_room(room_id, json.dumps({
                "type": "system", 
                "message": "User left.",
                "users": list(active_users[room_id])
            }))

async def broadcast_to_room(room_id: str, message: str):
    if room_id in active_rooms:
        for connection in list(active_rooms[room_id]):
            try:
                await connection.send_text(message)
            except:
                pass

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