app.py

import json
import asyncio
import shutil
from datetime import datetime, timedelta
from typing import Optional

import httpx
from fastapi import FastAPI, HTTPException, Depends, Request
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.templating import Jinja2Templates
from jose import JWTError, jwt
from pydantic import BaseModel

from config import settings
from database import (
    init_db, create_user, get_user_by_username, verify_password,
    get_user_by_api_key, create_api_key,
    get_user_api_keys, revoke_api_key, generate_api_key
)

app = FastAPI(title=settings.app_name)

security = HTTPBearer()
templates = Jinja2Templates(directory="templates")
app.mount("/static", StaticFiles(directory="static"), name="static")

init_db()

OLLAMA_MODEL = "llama3.1:8b"
OLLAMA_BASE = "http://localhost:11434"


# --- Models ---

class RegisterRequest(BaseModel):
    username: str
    email: str
    password: str

class LoginRequest(BaseModel):
    username: str
    password: str

class CreateApiKeyRequest(BaseModel):
    key_name: str

class ChatRequest(BaseModel):
    message: str
    stream: bool = True


# --- Auth Helpers ---

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    expire = datetime.utcnow() + (expires_delta or timedelta(minutes=settings.access_token_expire_minutes))
    to_encode.update({"exp": expire})
    return jwt.encode(to_encode, settings.secret_key, algorithm=settings.algorithm)

async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(security)):
    token = credentials.credentials
    if token.startswith("ollama_"):
        user = get_user_by_api_key(token)
        if not user:
            raise HTTPException(status_code=401, detail="Invalid API key")
        return user
    try:
        payload = jwt.decode(token, settings.secret_key, algorithms=[settings.algorithm])
        username = payload.get("sub")
        if username is None:
            raise HTTPException(status_code=401, detail="Invalid token")
    except JWTError:
        raise HTTPException(status_code=401, detail="Invalid token")
    user = get_user_by_username(username)
    if not user:
        raise HTTPException(status_code=401, detail="User not found")
    return user

async def get_user_or_none(request: Request):
    auth = request.headers.get("Authorization", "")
    if not auth.startswith("Bearer "):
        api_key = request.headers.get("X-API-Key", "")
        if api_key:
            return get_user_by_api_key(api_key)
        return None
    try:
        return await get_current_user(HTTPAuthorizationCredentials(scheme="Bearer", credentials=auth[7:]))
    except:
        return None


# --- Web Pages ---

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

@app.get("/login", response_class=HTMLResponse)
async def login_page(request: Request):
    return templates.TemplateResponse("login.html", {"request": request})

@app.get("/dashboard", response_class=HTMLResponse)
async def dashboard_page(request: Request):
    return templates.TemplateResponse("dashboard.html", {"request": request})

@app.get("/chat", response_class=HTMLResponse)
async def chat_page(request: Request):
    return templates.TemplateResponse("chat.html", {"request": request})


# --- Auth API ---

@app.post("/api/register")
async def register(req: RegisterRequest):
    user_id = create_user(req.username, req.email, req.password)
    if not user_id:
        raise HTTPException(status_code=400, detail="Username or email already taken")
    return {"message": "User created", "user_id": user_id}

@app.post("/api/login")
async def login(req: LoginRequest):
    user = get_user_by_username(req.username)
    if not user or not verify_password(req.password, user["password_hash"]):
        raise HTTPException(status_code=401, detail="Invalid credentials")
    token = create_access_token({"sub": user["username"]})
    return {"access_token": token, "token_type": "bearer", "username": user["username"]}

@app.get("/api/me")
async def me(user: dict = Depends(get_current_user)):
    return {
        "username": user["username"],
        "email": user["email"],
        "api_key": user["api_key"],
        "created_at": user["created_at"]
    }


# --- Model / Status ---

@app.get("/api/status")
async def status():
    ollama_path = shutil.which("ollama")
    if not ollama_path:
        return {"ollama": False, "model": False, "message": "Ollama not installed"}
    try:
        proc = await asyncio.create_subprocess_exec(ollama_path, "list", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
        stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=5)
        installed = [line.split()[0] for line in stdout.decode().strip().split("\n")[1:] if line.strip()]
        model_ready = OLLAMA_MODEL in installed
        return {
            "ollama": True,
            "model": model_ready,
            "installed": installed,
            "message": "Ready" if model_ready else f"Model {OLLAMA_MODEL} not yet installed (pulling in background)"
        }
    except Exception as e:
        return {"ollama": False, "model": False, "message": str(e)}


# --- API Key Management ---

@app.post("/api/api-keys")
async def create_new_api_key(req: CreateApiKeyRequest, user: dict = Depends(get_current_user)):
    key = create_api_key(user["id"], req.key_name)
    return {"api_key": key, "key_name": req.key_name}

@app.get("/api/api-keys")
async def list_api_keys(user: dict = Depends(get_current_user)):
    return {"api_keys": get_user_api_keys(user["id"])}

@app.delete("/api/api-keys/{key_id}")
async def delete_api_key(key_id: int, user: dict = Depends(get_current_user)):
    revoke_api_key(key_id, user["id"])
    return {"message": "API key revoked"}


# --- Chat API (streaming via SSE) ---

@app.post("/api/chat")
async def chat(req: ChatRequest, user: Optional[dict] = Depends(get_current_user)):
    if not user:
        raise HTTPException(status_code=401, detail="Authentication required")
    model_ready = await check_model_ready()
    if not model_ready:
        raise HTTPException(status_code=503, detail=f"{OLLAMA_MODEL} is still downloading. Please wait and try again.")

    async def generate():
        async with httpx.AsyncClient(timeout=None) as client:
            async with client.stream(
                "POST",
                f"{OLLAMA_BASE}/api/chat",
                json={"model": OLLAMA_MODEL, "messages": [{"role": "user", "content": req.message}], "stream": True},
            ) as resp:
                async for line in resp.aiter_lines():
                    if not line.strip():
                        continue
                    try:
                        chunk = json.loads(line)
                        content = chunk.get("message", {}).get("content", "")
                        if content:
                            yield f"data: {json.dumps({'content': content})}\n\n"
                        if chunk.get("done"):
                            yield f"data: {json.dumps({'done': True})}\n\n"
                    except json.JSONDecodeError:
                        continue

    return StreamingResponse(generate(), media_type="text/event-stream")

async def check_model_ready() -> bool:
    try:
        proc = await asyncio.create_subprocess_exec(
            shutil.which("ollama"), "list",
            stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
        )
        stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=5)
        return any(OLLAMA_MODEL in line for line in stdout.decode().strip().split("\n")[1:])
    except:
        return False


# --- Startup ---

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