main.py

from fastapi import FastAPI, UploadFile, File, HTTPException, Form, WebSocket
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, Response
from pydantic import BaseModel
import os
from typing import List, Optional
import json
import shutil

# Logic Imports
from logic.llm import LLMService
from logic.audio import AudioService
from logic.assets import AssetService

app = FastAPI(title="ConvoIQ ML Service")

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

# Services
llm_service = LLMService()
audio_service = AudioService()
asset_service = AssetService()

class SimulationRequest(BaseModel):
    avatar_id: str
    prompt: str

@app.get("/")
def read_root():
    return {"status": "healthy", "service": "ml-service", "gpu": "Available" if os.environ.get("CUDA_VISIBLE_DEVICES") else "CPU/Metal"}

@app.post("/generate-avatar")
async def generate_avatar(prompt: str = Form(...), image: UploadFile = File(None)):
    """
    Generates a 3D avatar from an image or text description.
    """
    if image:
        content = await image.read()
        # Use the Asset Service to generate real 3D model
        glb_path = await asset_service.generate_avatar_mesh(content)

        if glb_path and os.path.exists(glb_path):
            # Move to a static folder to serve
            os.makedirs("static/avatars", exist_ok=True)
            filename = f"avatar_{os.urandom(4).hex()}.glb"
            dest_path = f"static/avatars/{filename}"
            shutil.move(glb_path, dest_path)

            return {
                "status": "success",
                "avatar_url": f"http://localhost:8000/static/avatars/{filename}",
                "type": "generated_3d"
            }

    # Fallback/Mock for text-only or failure
    return {
        "status": "success",
        "avatar_url": "https://models.readyplayer.me/65e0b5b8e3c94c373786496c.glb",
        "note": "Used fallback avatar"
    }

@app.post("/simulate")
async def simulate_scene(request: SimulationRequest):
    """
    Main orchestration endpoint.
    1. Understands the prompt (LLM)
    2. Selects/Generates Environment (Splat)
    3. Selects/Generates Objects
    4. Returns Scene Graph
    """
    # 1. Reason about the scene
    scene_data = await llm_service.reason_about_scene(request.prompt)

    # 2. Get Environment
    env_desc = scene_data.get("environment_description", "A generic scene")
    env_url = await asset_service.generate_environment_splat(env_desc)

    # 3. Synthesize audio response
    response_text = scene_data.get("conversation_response", "I am starting the simulation now.")
    audio_data = await audio_service.synthesize(response_text)

    # Save audio to serve
    os.makedirs("static/audio", exist_ok=True)
    audio_filename = f"response_{os.urandom(4).hex()}.mp3"
    with open(f"static/audio/{audio_filename}", "wb") as f:
        f.write(audio_data)

    return {
        "status": "generated",
        "scene_config": scene_data,
        "environment_url": env_url,
        "audio_url": f"http://localhost:8000/static/audio/{audio_filename}",
        "subtitle": response_text
    }

@app.post("/transcribe")
async def transcribe_audio(file: UploadFile = File(...)):
    """
    Speech-to-Text endpoint.
    """
    content = await file.read()
    text = await audio_service.transcribe(content)
    return {"text": text}

@app.post("/chat")
async def chat(request: dict):
    """
    Text chat endpoint.
    """
    history = request.get("history", [])
    user_input = request.get("message", "")
    response = await llm_service.chat(history, user_input)
    return {"response": response}

# Serve static files
from fastapi.staticfiles import StaticFiles
os.makedirs("static", exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")

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