from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
import os

# =====================================================
# FastAPI App Setup
# =====================================================
app = FastAPI(title="AI Chat + Summarization + Vision API")

# =====================================================
# API Schemas
# =====================================================
class ChatRequest(BaseModel):
    message: str
    max_new_tokens: int = 80
    temperature: float = 0.7

class SummaryRequest(BaseModel):
    text: str
    max_length: int = 100
    min_length: int = 25

# =====================================================
# Endpoints (Stubbed for Testing)
# =====================================================
@app.post("/api/chat")
def chat_generate(req: ChatRequest):
    return {
        "success": True, 
        "response": f"This is a test response to: {req.message}",
        "tokens": [
            {"token": "This", "alternatives": [{"token": "That", "probability": 0.3}]},
            {"token": " is", "alternatives": [{"token": " was", "probability": 0.2}]},
            {"token": " a", "alternatives": [{"token": " the", "probability": 0.4}]}
        ]
    }

@app.post("/api/summarize")
def summarize_text(req: SummaryRequest):
    return {
        "success": True, 
        "summary": f"Summary: {req.text[:100]}..."
    }

@app.post("/predict_words")
def predict_words(req: dict):
    return [
        {"word": "test", "probability": 0.8},
        {"word": "demo", "probability": 0.1},
        {"word": "example", "probability": 0.1}
    ]

@app.get("/api/health")
def health_check():
    return {
        "status": "healthy", 
        "models": ["stub-models-for-testing"]
    }

# =====================================================
# Static Files
# =====================================================
if os.path.exists("static"):
    app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/")
def read_root():
    if os.path.exists("static/index.html"):
        return FileResponse("static/index.html")
    return {"message": "Minimal AI Chat API running!"}