from flask_socketio import SocketIO
from pymongo import MongoClient
from qdrant_client import QdrantClient
import logging
import torch
from sentence_transformers import SentenceTransformer
from datetime import datetime

from utils.config import Config

# Khởi tạo logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

# Khởi tạo các dịch vụ
class Extensions:
    def __init__(self):
        self.socketio = None
        self.mongo = None
        self.db = None
        self.qdrant_client = None
        self.model = None
        self.connected_clients = {}
    
    def get_current_time(self):
        """Get current UTC time"""
        return datetime.utcnow()
    
    def init_mongodb(self):
        try:
            self.mongo = MongoClient(
                Config.MONGO_URI,
                serverSelectionTimeoutMS=5000,
                connectTimeoutMS=10000,
                socketTimeoutMS=10000
            )
            self.mongo.admin.command('ping')
            self.db = self.mongo.get_database('legal_assistant')
            logging.info("MongoDB connection successful")
            return True
        except Exception as e:
            logging.error(f"MongoDB connection failed: {e}")
            try:
                self.mongo = MongoClient(
                    'mongodb://localhost:27017/',
                    serverSelectionTimeoutMS=5000
                )
                self.mongo.admin.command('ping')
                self.db = self.mongo.get_database('legal_assistant')
                logging.info("Local MongoDB connection successful")
                return True
            except Exception as e2:
                logging.error(f"Local MongoDB connection also failed: {e2}")
                return False
    
    def init_qdrant(self):
        try:
            self.qdrant_client = QdrantClient(
                url=Config.QDRANT_URL,
                api_key=Config.QDRANT_API_KEY,
            )
            logging.info("Qdrant client initialized")
            return True
        except Exception as e:
            logging.error(f"Failed to initialize Qdrant client: {e}")
            return False
    
    def init_embedding_model(self):
        try:
            self.model = SentenceTransformer(
                'hiieu/halong_embedding',
                device='cuda' if torch.cuda.is_available() else 'cpu'
            )
            logging.info("Embedding model initialized")
            return True
        except Exception as e:
            logging.error(f"Failed to initialize embedding model: {e}")
            return False
    
    def init_socketio(self, app):
        self.socketio = SocketIO(
            app,
            cors_allowed_origins="*",
            async_mode='threading',
            ping_timeout=120,
            ping_interval=30
        )
        logging.info("SocketIO initialized")
        return self.socketio
    
    def initialize_all(self, app):
        if not self.init_mongodb():
            return False
        if not self.init_qdrant():
            return False
        if not self.init_embedding_model():
            return False
        self.init_socketio(app)
        return True

# Tạo instance toàn cục
ext = Extensions()