# algoforge_prime/core/llm_clients.py
import os
import google.generativeai as genai
from huggingface_hub import InferenceClient

# --- Configuration ---
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
HF_TOKEN = os.getenv("HF_TOKEN")

GEMINI_API_CONFIGURED = False
HF_API_CONFIGURED = False

hf_inference_client = None
google_gemini_models = {} # To store initialized Gemini model instances

# --- Initialization ---
def initialize_clients():
    global GEMINI_API_CONFIGURED, HF_API_CONFIGURED, hf_inference_client, google_gemini_models
    
    # Google Gemini
    if GOOGLE_API_KEY:
        try:
            genai.configure(api_key=GOOGLE_API_KEY)
            GEMINI_API_CONFIGURED = True
            print("INFO: llm_clients.py - Google Gemini API configured.")
        except Exception as e:
            print(f"ERROR: llm_clients.py - Failed to configure Google Gemini API: {e}")
    else:
        print("WARNING: llm_clients.py - GOOGLE_API_KEY not found.")

    # Hugging Face
    if HF_TOKEN:
        try:
            hf_inference_client = InferenceClient(token=HF_TOKEN)
            HF_API_CONFIGURED = True
            print("INFO: llm_clients.py - Hugging Face InferenceClient initialized.")
        except Exception as e:
            print(f"ERROR: llm_clients.py - Failed to initialize Hugging Face InferenceClient: {e}")
    else:
        print("WARNING: llm_clients.py - HF_TOKEN not found.")

# Call initialize_clients when the module is imported for the first time.
# However, for Gradio apps that might reload, it's often better to call this explicitly from app.py's main scope.
# For now, let's assume it's called once. If you see issues, move the call.
# initialize_clients() # Or call this from app.py

def get_gemini_model_instance(model_id, system_instruction=None):
    """Gets or creates a Gemini model instance."""
    if not GEMINI_API_CONFIGURED:
        raise ConnectionError("Google Gemini API not configured.")
    
    instance_key = model_id + ("_sys" if system_instruction else "") # Simple keying
    if instance_key not in google_gemini_models:
        try:
            google_gemini_models[instance_key] = genai.GenerativeModel(
                model_name=model_id,
                system_instruction=system_instruction
            )
            print(f"INFO: Initialized Gemini Model Instance: {instance_key}")
        except Exception as e:
            print(f"ERROR: Failed to initialize Gemini model {model_id}: {e}")
            raise  # Re-raise the exception to be caught by the caller
    return google_gemini_models[instance_key]


class LLMResponse:
    def __init__(self, text=None, error=None, success=True, raw_response=None):
        self.text = text
        self.error = error
        self.success = success
        self.raw_response = raw_response # Store original API response if needed

    def __str__(self):
        if self.success:
            return self.text if self.text is not None else ""
        return f"ERROR: {self.error}"


def call_huggingface_api(prompt_text, model_id, temperature=0.7, max_new_tokens=350, system_prompt_text=None):
    if not HF_API_CONFIGURED or not hf_inference_client:
        return LLMResponse(error="Hugging Face API not configured.", success=False)

    full_prompt = prompt_text
    if system_prompt_text: # Simple prepend, specific formatting depends on model
        full_prompt = f"<s>[INST] <<SYS>>\n{system_prompt_text}\n<</SYS>>\n\n{prompt_text} [/INST]" # Llama-style

    try:
        use_sample = temperature > 0.0
        raw_response = hf_inference_client.text_generation(
            full_prompt, model=model_id, max_new_tokens=max_new_tokens,
            temperature=temperature if use_sample else None,
            do_sample=use_sample, stream=False
        )
        return LLMResponse(text=raw_response, raw_response=raw_response)
    except Exception as e:
        error_msg = f"HF API Error ({model_id}): {type(e).__name__} - {e}"
        print(f"ERROR: llm_clients.py - {error_msg}")
        return LLMResponse(error=error_msg, success=False, raw_response=e)

def call_gemini_api(prompt_text, model_id, temperature=0.7, max_new_tokens=400, system_prompt_text=None):
    if not GEMINI_API_CONFIGURED:
        return LLMResponse(error="Google Gemini API not configured.", success=False)

    try:
        model_instance = get_gemini_model_instance(model_id, system_instruction=system_prompt_text)
        
        generation_config = genai.types.GenerationConfig(
            temperature=temperature,
            max_output_tokens=max_new_tokens
        )
        raw_response = model_instance.generate_content(
            prompt_text, # User prompt
            generation_config=generation_config,
            stream=False
        )

        if raw_response.prompt_feedback and raw_response.prompt_feedback.block_reason:
            reason = raw_response.prompt_feedback.block_reason_message or raw_response.prompt_feedback.block_reason
            error_msg = f"Gemini API: Prompt blocked due to safety. Reason: {reason}"
            print(f"WARNING: llm_clients.py - {error_msg}")
            return LLMResponse(error=error_msg, success=False, raw_response=raw_response)

        if not raw_response.candidates or not raw_response.candidates[0].content.parts:
            finish_reason = raw_response.candidates[0].finish_reason if raw_response.candidates else "Unknown"
            if str(finish_reason).upper() == "SAFETY":
                error_msg = f"Gemini API: Response generation stopped by safety filters. Finish Reason: {finish_reason}"
            else:
                error_msg = f"Gemini API: Empty response or no content. Finish Reason: {finish_reason}"
            print(f"WARNING: llm_clients.py - {error_msg}")
            return LLMResponse(error=error_msg, success=False, raw_response=raw_response)
        
        return LLMResponse(text=raw_response.candidates[0].content.parts[0].text, raw_response=raw_response)

    except Exception as e:
        error_msg = f"Gemini API Error ({model_id}): {type(e).__name__} - {e}"
        print(f"ERROR: llm_clients.py - {error_msg}")
        return LLMResponse(error=error_msg, success=False, raw_response=e)