transcription/backend/main.py

"""FastAPI backend для сервиса транскрибации."""

import json
from pathlib import Path
from typing import List, Optional

from fastapi import FastAPI, File, UploadFile, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, PlainTextResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles

from backend.queue import (
    UPLOAD_DIR,
    PROCESSED_DIR,
    enqueue,
    get_all_tasks,
    get_task_status,
    get_processed_tree,
    read_file_content,
    set_progress_callback,
)

app = FastAPI(title="Transcription Service", version="1.0.0")

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

# WebSocket менеджер
class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []
    
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
    
    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
    
    async def broadcast(self, message: dict):
        for conn in self.active_connections:
            try:
                await conn.send_json(message)
            except Exception:
                pass

manager = ConnectionManager()

# Устанавливаем callback для отправки прогресса через WebSocket
set_progress_callback(manager.broadcast)


# === API Endpoints ===

@app.get("/", response_class=HTMLResponse)
async def root():
    """Главная страница."""
    index_path = Path(__file__).parent / "static" / "index.html"
    if index_path.exists():
        return index_path.read_text(encoding="utf-8")
    return "<h1>Transcription Service</h1><p>Frontend not built</p>"


@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    """Загружает файл и добавляет в очередь обработки."""
    # Сохраняем файл
    file_path = UPLOAD_DIR / file.filename
    with open(file_path, "wb") as f:
        content = await file.read()
        f.write(content)
    
    # Добавляем в очередь
    task_id = await enqueue(file_path)
    
    return {
        "task_id": task_id,
        "filename": file.filename,
        "status": "queued",
        "message": "Файл добавлен в очередь обработки",
    }


@app.post("/upload-batch")
async def upload_batch(files: List[UploadFile] = File(...)):
    """Загружает несколько файлов пакетно."""
    results = []
    for file in files:
        file_path = UPLOAD_DIR / file.filename
        with open(file_path, "wb") as f:
            content = await file.read()
            f.write(content)
        
        task_id = await enqueue(file_path)
        results.append({
            "task_id": task_id,
            "filename": file.filename,
            "status": "queued",
        })
    
    return {
        "uploaded": len(results),
        "tasks": results,
    }


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket для получения прогресса обработки."""
    await manager.connect(websocket)
    try:
        while True:
            # Ждём сообщения от клиента (ping/keepalive)
            data = await websocket.receive_text()
            msg = json.loads(data)
            
            if msg.get("action") == "get_tasks":
                tasks = get_all_tasks()
                await websocket.send_json({
                    "type": "tasks_list",
                    "tasks": tasks,
                })
            elif msg.get("action") == "get_tree":
                tree = get_processed_tree()
                await websocket.send_json({
                    "type": "file_tree",
                    "tree": tree,
                })
    except WebSocketDisconnect:
        manager.disconnect(websocket)
    except Exception:
        manager.disconnect(websocket)


@app.get("/api/tasks")
async def api_tasks():
    """Возвращает список всех задач."""
    return {"tasks": get_all_tasks()}


@app.get("/api/tasks/{task_id}")
async def api_task(task_id: str):
    """Возвращает статус конкретной задачи."""
    status = get_task_status(task_id)
    if not status:
        return {"error": "Task not found"}
    return status


@app.get("/api/files")
async def api_files():
    """Возвращает дерево обработанных файлов."""
    return {"tree": get_processed_tree()}


@app.get("/api/files/content")
async def api_file_content(path: str):
    """Возвращает содержимое файла."""
    try:
        content = read_file_content(path)
        return {"content": content, "path": path}
    except Exception as e:
        return {"error": str(e)}


@app.get("/api/files/download")
async def api_download(path: str):
    """Скачивает файл."""
    file_path = PROCESSED_DIR / path
    if not file_path.exists():
        return {"error": "File not found"}
    return FileResponse(file_path, filename=file_path.name)


# Статические файлы
app.mount("/static", StaticFiles(directory="backend/static"), name="static")