scud_ai/modules/web_api/routers/chat.py

"""
===============================================================================
FILE: modules/web_api/routers/chat.py
ROLE: Полнофункциональный роутер чата с извлечением текста из PDF и сканов,
      поддержкой Function Calling, Fast-Path и оптического распознавания OCR.
===============================================================================
"""

import os
import shutil
import base64
import logging
from typing import Optional, List, Dict, Any

from fastapi import APIRouter, Header, HTTPException, UploadFile, File, Form
from pydantic import BaseModel

from llm.agent import process_chat_message
from config import BASE_DIR

logger = logging.getLogger("CHAT_API")
router = APIRouter(prefix="/api/v1", tags=["Chat"])

UPLOAD_TMP_DIR = os.path.join(BASE_DIR, "data", "uploads")
os.makedirs(UPLOAD_TMP_DIR, exist_ok=True)


class ChatMessageRequest(BaseModel):
    message: str
    session_id: Optional[str] = "web_session_main"
    user_id: Optional[int] = 1


def resolve_user_id(authorization: Optional[str] = None, explicit_user_id: Optional[int] = None) -> int:
    if explicit_user_id and explicit_user_id > 0:
        return explicit_user_id

    if authorization and authorization.startswith("Bearer "):
        token = authorization.replace("Bearer ", "").strip()
        if token.isdigit():
            return int(token)
        elif token.startswith("dev_token_"):
            try:
                return int(token.replace("dev_token_", ""))
            except ValueError:
                pass
    return 1


@router.post("/chat")
async def chat_endpoint(payload: ChatMessageRequest, authorization: Optional[str] = Header(None)):
    user_id = resolve_user_id(authorization, payload.user_id)
    session_id = payload.session_id or "web_session_main"
    user_msg = payload.message.strip()

    if not user_msg:
        raise HTTPException(status_code=400, detail="Пустое сообщение")

    logger.info(f"Сообщение от user_id={user_id}, session_id={session_id}: {user_msg}")

    reply_text, history, action_payload = process_chat_message(
        user_id=user_id,
        user_message=user_msg,
        session_id=session_id
    )

    return {
        "status": "success",
        "user_id": user_id,
        "session_id": session_id,
        "response": reply_text,
        "action_payload": action_payload
    }


@router.post("/chat/upload")
async def chat_upload_endpoint(
    file: UploadFile = File(...),
    message: Optional[str] = Form(""),
    session_id: Optional[str] = Form("web_session_main"),
    authorization: Optional[str] = Header(None)
):
    user_id = resolve_user_id(authorization, 1)
    file_path = os.path.join(UPLOAD_TMP_DIR, file.filename)
    
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    file_context = ""
    image_b64 = None
    fn_lower = file.filename.lower()

    # 1. Текстовые форматы
    if fn_lower.endswith((".txt", ".csv", ".log", ".md")):
        try:
            with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
                file_context = f.read(6000)
        except Exception as e:
            logger.warning(f"Не удалось прочитать текст: {e}")

    # 2. Изображения (прямой OCR)
    elif fn_lower.endswith((".png", ".jpg", ".jpeg", ".webp")):
        try:
            with open(file_path, "rb") as f:
                image_b64 = base64.b64encode(f.read()).decode("utf-8")
        except Exception as e:
            logger.warning(f"Ошибка кодирования картинки в base64: {e}")

    # 3. PDF документы (текстовый слой + рендеринг скана при необходимости)
    elif fn_lower.endswith(".pdf"):
        # Попытка извлечь встроенный текстовый слой
        try:
            import pypdf
            reader = pypdf.PdfReader(file_path)
            extracted = []
            for page in reader.pages:
                t = page.extract_text()
                if t:
                    extracted.append(t)
            file_context = "\n".join(extracted).strip()
        except Exception:
            pass

        # Если текстового слоя мало (скан или фото документа), рендерим страницу в картинку для Vision OCR
        if len(file_context) < 40:
            try:
                import fitz  # PyMuPDF
                doc = fitz.open(file_path)
                if len(doc) > 0:
                    page = doc[0]
                    pix = page.get_pixmap(dpi=150)
                    img_bytes = pix.tobytes("png")
                    image_b64 = base64.b64encode(img_bytes).decode("utf-8")
                    file_context = ""
            except Exception as e:
                logger.warning(f"PyMuPDF не установлен или сбой рендеринга PDF: {e}")

    user_msg = message.strip() or f"Распознай и проанализируй прикрепленный документ {file.filename}"

    reply_text, history, action_payload = process_chat_message(
        user_id=user_id,
        user_message=user_msg,
        file_context=file_context,
        image_b64=image_b64,
        session_id=session_id
    )

    return {
        "status": "success",
        "user_id": user_id,
        "session_id": session_id,
        "response": reply_text,
        "action_payload": action_payload
    }