from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
from transformers import pipeline
import numpy as np
from typing import List
import asyncio
from concurrent.futures import ThreadPoolExecutor

app = FastAPI()

# Создаем пул потоков для фоновых задач
executor = ThreadPoolExecutor(max_workers=2)

# Загружаем модели (один раз при старте)
print("Загрузка моделей...")
sentence_model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
sentiment_pipeline = pipeline("sentiment-analysis", model="blanchefort/rubert-base-cased-sentiment")
print("Модели загружены!")

class TextRequest(BaseModel):
    text: str

class EmbeddingResponse(BaseModel):
    embedding: List[float]

class SentimentResponse(BaseModel):
    label: str
    score: float

class TextsRequest(BaseModel):
    texts: List[str]

class SimilarityRequest(BaseModel):
    text1: str
    text2: str

class SimilarityResponse(BaseModel):
    similarity: float

@app.get("/")
def root():
    return {"message": "AI Service for Grant Platform", "status": "running"}

@app.get("/health")
def health():
    return {"status": "ok", "models_loaded": True}

@app.post("/embed", response_model=EmbeddingResponse)
def get_embedding(request: TextRequest):
    """Возвращает эмбеддинг текста"""
    try:
        embedding = sentence_model.encode(request.text)
        embedding_list = embedding.tolist() if isinstance(embedding, np.ndarray) else embedding
        return EmbeddingResponse(embedding=embedding_list)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/sentiment", response_model=SentimentResponse)
def get_sentiment(request: TextRequest):
    """Анализ тональности текста"""
    try:
        result = sentiment_pipeline(request.text[:512])[0]
        return SentimentResponse(label=result['label'], score=result['score'])
    except Exception as e:
        return SentimentResponse(label="NEUTRAL", score=0.5)

@app.post("/batch_embed")
def batch_embed(request: TextsRequest):
    """Массовое получение эмбеддингов для нескольких текстов"""
    try:
        embeddings = sentence_model.encode(request.texts)
        return {"embeddings": [e.tolist() for e in embeddings]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/similarity", response_model=SimilarityResponse)
def get_similarity(request: SimilarityRequest):
    """Косинусное сходство между двумя текстами"""
    try:
        emb1 = sentence_model.encode(request.text1)
        emb2 = sentence_model.encode(request.text2)
        similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
        return SimilarityResponse(similarity=float(similarity))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))