#!/usr/bin/env python3
"""
Standalone RAG сервер для демонстрации
Работает напрямую с Qdrant и OpenRouter API без микросервисов
"""

from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
import uvicorn
import os
from pathlib import Path
import uuid
from datetime import datetime
import httpx

from qdrant_client import QdrantClient
from qdrant_client.models import PointStruct
from sentence_transformers import SentenceTransformer

# Конфигурация
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")  # Для Qdrant Cloud
COLLECTION_NAME = "onboarding_documents"

# OpenRouter Configuration
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "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")
OPENROUTER_MODEL = os.getenv("OPENROUTER_MODEL", "mistralai/voxtral-small-24b-2507")
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1/chat/completions"

# Redis (опционально, для кэширования)
REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
REDIS_PORT = int(os.getenv("REDIS_PORT", "6379"))
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD")

# Инициализация
app = FastAPI(title="AI-RAG Onboarding Demo", version="1.0.0")

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

@app.get("/")
async def root():
    """Главная страница"""
    return {
        "service": "AI-RAG Onboarding Assistant",
        "status": "running",
        "version": "1.0.0",
        "endpoints": {
            "health": "/health",
            "query": "/api/v1/query",
            "upload": "/api/v1/upload",
            "stats": "/api/v1/stats"
        }
    }


# Глобальные объекты
qdrant_client = None
embedding_model = None
httpx_client = None


class QueryRequest(BaseModel):
    query: str
    dept_id: str = "onboarding"
    user_id: str = "demo_user"
    session_id: Optional[str] = None


class Source(BaseModel):
    text: str
    score: float
    metadata: Dict[str, Any]


class QueryResponse(BaseModel):
    answer: str
    sources: List[Source]
    metadata: Dict[str, Any]


class DocumentUploadResponse(BaseModel):
    document_id: str
    chunks_created: int
    message: str
    metadata: Dict[str, Any]


@app.on_event("startup")
async def startup_event():
    """Инициализация при запуске"""
    global qdrant_client, embedding_model, httpx_client
    
    print("🚀 Запуск RAG сервера...")
    
    # Qdrant
    print("🔌 Подключение к Qdrant...")
    # In-memory Qdrant (для HF Spaces)
    from qdrant_client.models import Distance, VectorParams
    qdrant_client = QdrantClient(":memory:")
    print("✅ Qdrant подключен (in-memory режим)")
    
    # Создаем коллекцию если её нет
    try:
        qdrant_client.get_collection(COLLECTION_NAME)
        print(f"✅ Коллекция {COLLECTION_NAME} уже существует")
    except:
        qdrant_client.create_collection(
            collection_name=COLLECTION_NAME,
            vectors_config=VectorParams(size=384, distance=Distance.COSINE)
        )
        print(f"✅ Коллекция {COLLECTION_NAME} создана (384 dimensions)")
    
    # Embedding model
    print("🧠 Загрузка модели эмбеддингов...")
    # Используем многоязычную модель для правильной работы с русским языком
    embedding_model = SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')
    print("✅ Модель эмбеддингов загружена (многоязычная)")
    
    # HTTP client для OpenRouter
    print("🤖 Настройка OpenRouter API...")
    httpx_client = httpx.AsyncClient(timeout=15.0)
    print(f"✅ OpenRouter настроен (модель: {OPENROUTER_MODEL})")
    
    print("✨ Сервер готов!")


@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "qdrant": "connected" if qdrant_client else "disconnected",
        "embedding_model": "loaded" if embedding_model else "not loaded",
        "llm": "configured" if httpx_client else "not configured"
    }


@app.post("/api/v1/query", response_model=QueryResponse)
async def query(request: QueryRequest):
    """Основной endpoint для запросов"""
    
    try:
        import time
        start_time = time.time()
        
        # 1. Генерация эмбеддинга запроса
        query_embedding = embedding_model.encode(request.query).tolist()
        
        # 2. Поиск в Qdrant с ограничением количества
        # Используем только limit без score_threshold чтобы получить топ-3 САМЫХ релевантных
        # даже если их score не очень высокий
        search_results = qdrant_client.search(
            collection_name=COLLECTION_NAME,
            query_vector=query_embedding,
            limit=3,  # Только топ-3 самых релевантных по score
            with_payload=True
        )
        
        # 3. Подготовка контекста
        sources = []
        context_parts = []
        
        for idx, hit in enumerate(search_results, 1):
            source = Source(
                text=hit.payload.get('text', ''),
                score=hit.score,
                metadata={
                    'title': hit.payload.get('title', 'Unknown'),
                    'chunk_index': hit.payload.get('chunk_index', 0),
                    'doc_type': hit.payload.get('doc_type', 'unknown'),
                    'department': hit.payload.get('department', 'unknown'),
                    'last_updated': hit.payload.get('last_updated', 'unknown'),
                }
            )
            sources.append(source)
            context_parts.append(f"[Источник {idx}] {hit.payload.get('text', '')}")
        
        context = "\n\n".join(context_parts)
        
        # 4. Генерация ответа с помощью Gemini
        prompt = f"""Ты - помощник по онбордингу новых сотрудников компании. Используй предоставленный контекст для ответа на вопрос.

КОНТЕКСТ:
{context}

ВОПРОС: {request.query}

ИНСТРУКЦИИ:
- Отвечай на русском языке используя Markdown форматирование
- Используй только информацию из контекста
- Структурируй ответ понятно и кратко с помощью заголовков (##, ###), списков (-, *), жирного текста (**текст**)
- **ВАЖНО ДЛЯ ССЫЛОК:** Оборачивай ЦЕЛЫЕ ФРАЗЫ в ссылки, а не отдельные слова. Ссылка должна читаться естественно как часть предложения.
  
  ✅ ПРАВИЛЬНО (естественное чтение):
  - "Встреча в [кабинете 101](source:1)" 
  - "Обратитесь к [HR-менеджеру](source:2)"
  - "Перейдите на [https://account.company.kz](source:3)"
  - "Временный пароль действителен [24 часа](source:1)"
  - "Получите логин от [IT-отдела](source:2)"
  
  ❌ НЕПРАВИЛЬНО (разрывает текст):
  - "Временный пароль от [HR-менеджера](source:1) для входа" (разрывает фразу)
  - "Действителен только [24 часа](source:1) !" (отдельные слова)
  
- Делай ссылки МИНИМАЛЬНЫМИ - только ключевой факт, не целое предложение
- Используй ссылки для конкретных фактов: адреса сайтов, номера кабинетов, имена должностей, временные интервалы
- Если информации недостаточно, честно скажи об этом
- Будь дружелюбным и помогающим

ОТВЕТ:"""

        # 4. Генерация ответа с помощью OpenRouter
        try:
            response = await httpx_client.post(
                OPENROUTER_BASE_URL,
                headers={
                    "Authorization": f"Bearer {OPENROUTER_API_KEY}",
                    "HTTP-Referer": "https://github.com/baltabekpro/ai-rag-core",
                    "X-Title": "AI-RAG Onboarding"
                },
                json={
                    "model": OPENROUTER_MODEL,
                    "messages": [
                        {
                            "role": "user",
                            "content": prompt
                        }
                    ],
                    "temperature": 0.3,
                    "max_tokens": 800
                },
                timeout=15.0
            )
            response.raise_for_status()
            
            result = response.json()
            answer = result["choices"][0]["message"]["content"]
            
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 429:
                raise HTTPException(
                    status_code=503,
                    detail="AI сервис временно перегружен. Попробуйте через минуту."
                )
            raise HTTPException(
                status_code=500,
                detail=f"Ошибка OpenRouter API: {e.response.text}"
            )
        except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"Ошибка генерации ответа: {str(e)}"
            )
        
        # 5. Метаданные
        processing_time = int((time.time() - start_time) * 1000)
        
        return QueryResponse(
            answer=answer,
            sources=sources,
            metadata={
                "processing_time": processing_time,
                "model": OPENROUTER_MODEL,
                "sources_count": len(sources),
                "department": request.dept_id
            }
        )
        
    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        print(f"❌ ОШИБКА в /api/v1/query: {error_trace}")
        raise HTTPException(status_code=500, detail=f"Ошибка обработки запроса: {str(e)}")


@app.get("/api/v1/stats")
async def get_stats():
    """Статистика системы"""
    try:
        collection_info = qdrant_client.get_collection(COLLECTION_NAME)
        return {
            "collection": COLLECTION_NAME,
            "documents_count": collection_info.points_count,
            "vector_size": collection_info.config.params.vectors.size,
            "status": "operational"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/v1/documents/upload", response_model=DocumentUploadResponse)
async def upload_document(
    file: UploadFile = File(...),
    title: Optional[str] = Form(None),
    department: str = Form("onboarding"),
    doc_type: str = Form("guide")
):
    """
    Загрузка и индексация документа
    
    Поддерживаемые форматы: .txt, .md
    
    Процесс:
    1. Чтение файла
    2. Разбивка на чанки (512 токенов)
    3. Генерация эмбеддингов
    4. Сохранение в Qdrant
    """
    try:
        import time
        start_time = time.time()
        
        # Проверка формата файла
        allowed_extensions = ['.txt', '.md']
        file_ext = os.path.splitext(file.filename)[1].lower()
        
        if file_ext not in allowed_extensions:
            raise HTTPException(
                status_code=400, 
                detail=f"Неподдерживаемый формат файла. Разрешены: {', '.join(allowed_extensions)}"
            )
        
        # Чтение содержимого
        content = await file.read()
        text = content.decode('utf-8')
        
        if not text.strip():
            raise HTTPException(status_code=400, detail="Файл пустой")
        
        # Используем название файла если title не указан
        doc_title = title or file.filename
        
        # ID документа
        document_id = str(uuid.uuid4())
        
        # Разбивка на чанки (простая - по 512 токенов ~2000 символов)
        chunk_size = 2000
        overlap = 200
        chunks = []
        
        for i in range(0, len(text), chunk_size - overlap):
            chunk = text[i:i + chunk_size]
            if chunk.strip():
                chunks.append(chunk.strip())
        
        if not chunks:
            raise HTTPException(status_code=400, detail="Не удалось создать чанки из документа")
        
        # Генерация эмбеддингов и сохранение
        points = []
        
        for idx, chunk in enumerate(chunks):
            # Эмбеддинг
            embedding = embedding_model.encode(chunk).tolist()
            
            # Point для Qdrant
            point = PointStruct(
                id=str(uuid.uuid4()),
                vector=embedding,
                payload={
                    "text": chunk,
                    "document_id": document_id,
                    "chunk_index": idx,
                    "title": doc_title,
                    "department": department,
                    "doc_type": doc_type,
                    "last_updated": datetime.utcnow().isoformat(),
                    "filename": file.filename
                }
            )
            points.append(point)
        
        # Загрузка в Qdrant
        qdrant_client.upsert(
            collection_name=COLLECTION_NAME,
            points=points
        )
        
        processing_time = int((time.time() - start_time) * 1000)
        
        return DocumentUploadResponse(
            document_id=document_id,
            chunks_created=len(chunks),
            message=f"Документ '{doc_title}' успешно загружен и проиндексирован",
            metadata={
                "processing_time_ms": processing_time,
                "filename": file.filename,
                "file_size": len(content),
                "chunk_size": chunk_size,
                "department": department,
                "doc_type": doc_type
            }
        )
        
    except HTTPException:
        raise
    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        print(f"❌ ОШИБКА в /api/v1/documents/upload: {error_trace}")
        raise HTTPException(status_code=500, detail=f"Ошибка загрузки документа: {str(e)}")


if __name__ == "__main__":
    print("=" * 70)
    print("  AI-RAG Onboarding - Standalone Demo Server")
    print("=" * 70)
    print()
    print("📍 Сервер запускается на: http://localhost:8081")
    print("📄 API документация: http://localhost:8081/docs")
    print("💬 Откройте frontend/chat.html и измените API URL на http://localhost:8081")
    print()
    
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8081,
        log_level="info"
    )