# storyverse_weaver/core/llm_services.py
import os
import google.generativeai as genai
from huggingface_hub import InferenceClient
# from dotenv import load_dotenv # Optional: for local .env file
# load_dotenv() # Load environment variables from .env file if present

GOOGLE_API_KEY = os.getenv("STORYVERSE_GOOGLE_API_KEY") # Use specific env var names
HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN")

GEMINI_TEXT_CONFIGURED = False
HF_TEXT_CONFIGURED = False
hf_inference_text_client = None

class LLMTextResponse:
    def __init__(self, text=None, error=None, success=True, model_id_used="unknown_text_llm"):
        self.text, self.error, self.success, self.model_id_used = text, error, success, model_id_used

def initialize_text_llms():
    global GOOGLE_API_KEY, HF_TOKEN, GEMINI_TEXT_CONFIGURED, HF_TEXT_CONFIGURED, hf_inference_text_client
    print("INFO: llm_services.py - Initializing Text LLM clients...")
    if GOOGLE_API_KEY and GOOGLE_API_KEY.strip():
        try:
            genai.configure(api_key=GOOGLE_API_KEY)
            GEMINI_TEXT_CONFIGURED = True
            print("SUCCESS: llm_services.py - Google Gemini API (for text) configured.")
        except Exception as e:
            print(f"ERROR: llm_services.py - Failed to configure Google Gemini API: {e}")
            GEMINI_TEXT_CONFIGURED = False
    else:
        print("WARNING: llm_services.py - STORYVERSE_GOOGLE_API_KEY not found or empty.")
        GEMINI_TEXT_CONFIGURED = False

    if HF_TOKEN and HF_TOKEN.strip():
        try:
            hf_inference_text_client = InferenceClient(token=HF_TOKEN)
            HF_TEXT_CONFIGURED = True
            print("SUCCESS: llm_services.py - Hugging Face InferenceClient (for text) initialized.")
        except Exception as e:
            print(f"ERROR: llm_services.py - Failed to initialize HF InferenceClient: {e}")
            HF_TEXT_CONFIGURED = False
    else:
        print("WARNING: llm_services.py - STORYVERSE_HF_TOKEN not found or empty.")
        HF_TEXT_CONFIGURED = False
    print(f"INFO: llm_services.py - Text LLM Init complete. Gemini Text: {GEMINI_TEXT_CONFIGURED}, HF Text: {HF_TEXT_CONFIGURED}")

def is_gemini_text_ready(): return GEMINI_TEXT_CONFIGURED
def is_hf_text_ready(): return HF_TEXT_CONFIGURED

def generate_text_gemini(prompt: str, model_id: str = "gemini-1.5-flash-latest", system_prompt: str = None, temperature: float = 0.7, max_tokens: int = 512) -> LLMTextResponse:
    if not is_gemini_text_ready():
        return LLMTextResponse(error="Gemini text API not configured.", success=False, model_id_used=model_id)
    try:
        model = genai.GenerativeModel(model_name=model_id, system_instruction=system_prompt)
        config = genai.types.GenerationConfig(temperature=temperature, max_output_tokens=max_tokens)
        response = model.generate_content(prompt, generation_config=config)
        # Add robust response checking as in AlgoForge's llm_clients.py
        if response.prompt_feedback and response.prompt_feedback.block_reason:
            return LLMTextResponse(error=f"Gemini: Prompt blocked ({response.prompt_feedback.block_reason})", success=False, model_id_used=model_id)
        if not response.candidates or not response.candidates[0].content.parts:
            return LLMTextResponse(error=f"Gemini: No content generated (Finish reason: {response.candidates[0].finish_reason if response.candidates else 'Unknown'})", success=False, model_id_used=model_id)
        return LLMTextResponse(text=response.text, model_id_used=model_id)
    except Exception as e:
        return LLMTextResponse(error=f"Gemini API Error ({model_id}): {type(e).__name__} - {str(e)}", success=False, model_id_used=model_id)

def generate_text_hf(prompt: str, model_id: str = "mistralai/Mistral-7B-Instruct-v0.2", system_prompt: str = None, temperature: float = 0.7, max_tokens: int = 512) -> LLMTextResponse:
    if not is_hf_text_ready() or not hf_inference_text_client:
        return LLMTextResponse(error="HF text API not configured.", success=False, model_id_used=model_id)
    
    full_prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n{prompt} [/INST]" if system_prompt else prompt
    try:
        use_sample = temperature > 0.001
        response_text = hf_inference_text_client.text_generation(
            full_prompt, model=model_id, max_new_tokens=max_tokens,
            temperature=temperature if use_sample else None, do_sample=use_sample
        )
        return LLMTextResponse(text=response_text, model_id_used=model_id)
    except Exception as e:
        return LLMTextResponse(error=f"HF API Error ({model_id}): {type(e).__name__} - {str(e)}", success=False, model_id_used=model_id)

print("DEBUG: core.llm_services (for StoryVerseWeaver) - Module defined.")# storyverse_weaver/core/llm_services.py
import os
import google.generativeai as genai
from huggingface_hub import InferenceClient
# from dotenv import load_dotenv # Optional: for local .env file
# load_dotenv() # Load environment variables from .env file if present

GOOGLE_API_KEY = os.getenv("STORYVERSE_GOOGLE_API_KEY") # Use specific env var names
HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN")

GEMINI_TEXT_CONFIGURED = False
HF_TEXT_CONFIGURED = False
hf_inference_text_client = None

class LLMTextResponse:
    def __init__(self, text=None, error=None, success=True, model_id_used="unknown_text_llm"):
        self.text, self.error, self.success, self.model_id_used = text, error, success, model_id_used

def initialize_text_llms():
    global GOOGLE_API_KEY, HF_TOKEN, GEMINI_TEXT_CONFIGURED, HF_TEXT_CONFIGURED, hf_inference_text_client
    print("INFO: llm_services.py - Initializing Text LLM clients...")
    if GOOGLE_API_KEY and GOOGLE_API_KEY.strip():
        try:
            genai.configure(api_key=GOOGLE_API_KEY)
            GEMINI_TEXT_CONFIGURED = True
            print("SUCCESS: llm_services.py - Google Gemini API (for text) configured.")
        except Exception as e:
            print(f"ERROR: llm_services.py - Failed to configure Google Gemini API: {e}")
            GEMINI_TEXT_CONFIGURED = False
    else:
        print("WARNING: llm_services.py - STORYVERSE_GOOGLE_API_KEY not found or empty.")
        GEMINI_TEXT_CONFIGURED = False

    if HF_TOKEN and HF_TOKEN.strip():
        try:
            hf_inference_text_client = InferenceClient(token=HF_TOKEN)
            HF_TEXT_CONFIGURED = True
            print("SUCCESS: llm_services.py - Hugging Face InferenceClient (for text) initialized.")
        except Exception as e:
            print(f"ERROR: llm_services.py - Failed to initialize HF InferenceClient: {e}")
            HF_TEXT_CONFIGURED = False
    else:
        print("WARNING: llm_services.py - STORYVERSE_HF_TOKEN not found or empty.")
        HF_TEXT_CONFIGURED = False
    print(f"INFO: llm_services.py - Text LLM Init complete. Gemini Text: {GEMINI_TEXT_CONFIGURED}, HF Text: {HF_TEXT_CONFIGURED}")

def is_gemini_text_ready(): return GEMINI_TEXT_CONFIGURED
def is_hf_text_ready(): return HF_TEXT_CONFIGURED

def generate_text_gemini(prompt: str, model_id: str = "gemini-1.5-flash-latest", system_prompt: str = None, temperature: float = 0.7, max_tokens: int = 512) -> LLMTextResponse:
    if not is_gemini_text_ready():
        return LLMTextResponse(error="Gemini text API not configured.", success=False, model_id_used=model_id)
    try:
        model = genai.GenerativeModel(model_name=model_id, system_instruction=system_prompt)
        config = genai.types.GenerationConfig(temperature=temperature, max_output_tokens=max_tokens)
        response = model.generate_content(prompt, generation_config=config)
        # Add robust response checking as in AlgoForge's llm_clients.py
        if response.prompt_feedback and response.prompt_feedback.block_reason:
            return LLMTextResponse(error=f"Gemini: Prompt blocked ({response.prompt_feedback.block_reason})", success=False, model_id_used=model_id)
        if not response.candidates or not response.candidates[0].content.parts:
            return LLMTextResponse(error=f"Gemini: No content generated (Finish reason: {response.candidates[0].finish_reason if response.candidates else 'Unknown'})", success=False, model_id_used=model_id)
        return LLMTextResponse(text=response.text, model_id_used=model_id)
    except Exception as e:
        return LLMTextResponse(error=f"Gemini API Error ({model_id}): {type(e).__name__} - {str(e)}", success=False, model_id_used=model_id)

def generate_text_hf(prompt: str, model_id: str = "mistralai/Mistral-7B-Instruct-v0.2", system_prompt: str = None, temperature: float = 0.7, max_tokens: int = 512) -> LLMTextResponse:
    if not is_hf_text_ready() or not hf_inference_text_client:
        return LLMTextResponse(error="HF text API not configured.", success=False, model_id_used=model_id)
    
    full_prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n{prompt} [/INST]" if system_prompt else prompt
    try:
        use_sample = temperature > 0.001
        response_text = hf_inference_text_client.text_generation(
            full_prompt, model=model_id, max_new_tokens=max_tokens,
            temperature=temperature if use_sample else None, do_sample=use_sample
        )
        return LLMTextResponse(text=response_text, model_id_used=model_id)
    except Exception as e:
        return LLMTextResponse(error=f"HF API Error ({model_id}): {type(e).__name__} - {str(e)}", success=False, model_id_used=model_id)

print("DEBUG: core.llm_services (for StoryVerseWeaver) - Module defined.")