connections.py

""" 
Manages connections to external services: SSH, Solr, and Google Gemini.

This module centralizes the initialization logic, making the main application
cleaner and more focused on its primary task. It provides a single function
to set up all necessary connections.
"""


import pysolr
import google.generativeai as genai
from sshtunnel import SSHTunnelForwarder
import config

def initialize_connections():
    """
    Establishes the SSH tunnel, and initializes Solr and Gemini clients.

    Returns:
        A tuple containing the initialized (ssh_tunnel_server, solr_client, llm_model).
        Returns (None, None, None) if any part of the initialization fails.
    """
    ssh_tunnel_server = None
    try:
        # 1. Configure and start the SSH Tunnel
        ssh_tunnel_server = SSHTunnelForwarder(
            (config.SSH_HOST, config.SSH_PORT),
            ssh_username=config.SSH_USER,
            ssh_password=config.SSH_PASS,
            remote_bind_address=(config.REMOTE_SOLR_HOST, config.REMOTE_SOLR_PORT),
            local_bind_address=('127.0.0.1', config.LOCAL_BIND_PORT)
        )
        ssh_tunnel_server.start()
        print(f"🚀 SSH tunnel established: Local Port {ssh_tunnel_server.local_bind_port} -> Remote Solr.")

        # 2. Initialize the pysolr client
        solr_url = f'http://127.0.0.1:{ssh_tunnel_server.local_bind_port}/solr/{config.SOLR_CORE_NAME}'
        solr_client = pysolr.Solr(solr_url, auth=(config.SOLR_USER, config.SOLR_PASS), always_commit=True)
        solr_client.ping()
        print(f"✅ Solr connection successful on core '{config.SOLR_CORE_NAME}'.")

        # 3. Initialize the LLM
        genai.configure(api_key=config.GEMINI_API_KEY)
        llm_model = genai.GenerativeModel('gemini-2.5-flash', generation_config=genai.types.GenerationConfig(temperature=0, max_output_tokens=1048576))
        print(f"✅ LLM Model '{llm_model.model_name}' initialized.")

        print("✅ System Initialized Successfully.")
        return ssh_tunnel_server, solr_client, llm_model

    except Exception as e:
        print(f"\n❌ An error occurred during setup: {e}")
        if ssh_tunnel_server and ssh_tunnel_server.is_active:
            ssh_tunnel_server.stop()
        return None, None, None