agents/text_analyzer_agent.py

import os
import certifi
import logging
import subprocess # For calling ffmpeg if needed
from typing import List, Dict, Optional
from dotenv import load_dotenv

from llama_index.core.agent.workflow import ReActAgent
from llama_index.core.tools import FunctionTool
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import Document

# Attempt to import Whisper
try:
    import whisper
    WHISPER_AVAILABLE = True
except ImportError:
    logging.warning("openai-whisper not installed. Audio transcription tool will be unavailable.")
    WHISPER_AVAILABLE = False

# Load environment variables
load_dotenv()

# Setup logging
logger = logging.getLogger(__name__)

# Global Whisper model instance (lazy loaded)
_whisper_model = None

os.environ["SSL_CERT_FILE"] = certifi.where()

# Helper function to load prompt from file
def load_prompt_from_file(filename: str, default_prompt: str) -> str:
    """Loads a prompt from a text file."""
    try:
        script_dir = os.path.dirname(__file__)
        prompt_path = os.path.join(script_dir, filename)
        with open(prompt_path, "r") as f:
            prompt = f.read()
            logger.info(f"Successfully loaded prompt from {prompt_path}")
            return prompt
    except FileNotFoundError:
        logger.warning(f"Prompt file {filename} not found at {prompt_path}. Using default.")
        return default_prompt
    except Exception as e:
        logger.error(f"Error loading prompt file {filename}: {e}", exc_info=True)
        return default_prompt

# --- Helper function to load Whisper model ---
def _load_whisper_model(model_size: str = "small") -> Optional[object]:
    """Loads the Whisper model instance, lazy loading."""
    global _whisper_model
    if not WHISPER_AVAILABLE:
        logger.error("Whisper library not available, cannot load model.")
        return None
        
    if _whisper_model is None:
        try:
            logger.info(f"Loading Whisper model: {model_size}...")
            # Allow model size selection via env var, default to "base"
            selected_model_size = os.getenv("WHISPER_MODEL_SIZE", model_size)
            print(f"Available Whisper models: {whisper.available_models()}")
            _whisper_model = whisper.load_model(selected_model_size)
            logger.info(f"Whisper model {selected_model_size} loaded successfully.")
        except Exception as e:
            logger.error(f"Failed to load Whisper model {selected_model_size}: {e}", exc_info=True)
            _whisper_model = None # Ensure it remains None on failure
            
    return _whisper_model

# --- Tool Functions ---

def summarize_text(text: str, max_length: int = 150, min_length: int = 30) -> str:
    """Summarize the provided text using an LLM."""
    logger.info(f"Summarizing text (length: {len(text)} chars). Max/Min length: {max_length}/{min_length}")
    
    # Configuration for summarization LLM
    summarizer_llm_model = os.getenv("SUMMARIZER_LLM_MODEL", "models/gemini-1.5-flash") # Use flash for speed
    gemini_api_key = os.getenv("GEMINI_API_KEY")
    if not gemini_api_key:
        logger.error("GEMINI_API_KEY not found for summarization tool LLM.")
        return "Error: GEMINI_API_KEY not set for summarization."

    # Truncate input text if excessively long to avoid API limits/costs
    max_input_chars = 30000 # Example limit, adjust as needed
    if len(text) > max_input_chars:
        logger.warning(f"Input text truncated to {max_input_chars} chars for summarization.")
        text = text[:max_input_chars]

    prompt = (
        f"Summarize the following text concisely. Aim for a length between {min_length} and {max_length} words. "
        f"Focus on the main points and key information.\n\n"
        f"TEXT:\n{text}\n\nSUMMARY:"
    )

    try:
        llm = GoogleGenAI(api_key=gemini_api_key, model=summarizer_llm_model)
        logger.info(f"Using summarization LLM: {summarizer_llm_model}")
        response = llm.complete(prompt)
        summary = response.text.strip()
        logger.info(f"Summarization successful (output length: {len(summary.split())} words).")
        return summary
    except Exception as e:
        logger.error(f"LLM call failed during summarization: {e}", exc_info=True)
        return f"Error during summarization: {e}"

def extract_entities(text: str, entity_types: List[str] = ["PERSON", "ORG", "GPE", "DATE", "EVENT"]) -> Dict[str, List[str]]:
    """Extract named entities (like people, organizations, locations, dates) from the text using an LLM."""
    logger.info(f"Extracting entities (types: {entity_types}) from text (length: {len(text)} chars).")
    
    # Configuration for entity extraction LLM
    entity_llm_model = os.getenv("ENTITY_LLM_MODEL", "models/gemini-1.5-flash") # Use flash for speed
    gemini_api_key = os.getenv("GEMINI_API_KEY")
    if not gemini_api_key:
        logger.error("GEMINI_API_KEY not found for entity extraction tool LLM.")
        return {"error": "GEMINI_API_KEY not set for entity extraction."}

    # Truncate input text if excessively long
    max_input_chars = 30000 # Example limit
    if len(text) > max_input_chars:
        logger.warning(f"Input text truncated to {max_input_chars} chars for entity extraction.")
        text = text[:max_input_chars]

    # Define the desired output format clearly in the prompt
    prompt = (
        f"Extract named entities from the following text. Identify entities of these types: {', '.join(entity_types)}. "
        f"Format the output as a JSON object where keys are the entity types (uppercase) and values are lists of unique strings found for that type. "
        f"If no entities of a type are found, include the key with an empty list.\n\n"
        f"TEXT:\n{text}\n\nJSON_OUTPUT:"
    )

    try:
        llm = GoogleGenAI(api_key=gemini_api_key, model=entity_llm_model, response_mime_type="application/json") # Request JSON output
        logger.info(f"Using entity extraction LLM: {entity_llm_model}")
        response = llm.complete(prompt)
        
        # Attempt to parse the JSON response
        import json
        try:
            # The response might be wrapped in ```json ... ```, try to extract it
            json_str = response.text.strip()
            if json_str.startswith("```json"):
                json_str = json_str[7:]
            if json_str.endswith("```"):
                json_str = json_str[:-3]
            
            entities = json.loads(json_str.strip())
            # Validate structure (optional but good practice)
            if not isinstance(entities, dict):
                 raise ValueError("LLM response is not a JSON object.")
            # Ensure all requested types are present, even if empty
            for entity_type in entity_types:
                if entity_type not in entities:
                    entities[entity_type] = []
                elif not isinstance(entities[entity_type], list):
                     logger.warning(f"Entity type {entity_type} value is not a list, converting.")
                     entities[entity_type] = [str(entities[entity_type])] # Attempt conversion
                     
            logger.info(f"Entity extraction successful. Found entities: { {k: len(v) for k, v in entities.items()} }")
            return entities
        except json.JSONDecodeError as json_err:
            logger.error(f"Failed to parse JSON response from LLM: {json_err}. Response text: {response.text}")
            return {"error": f"Failed to parse LLM JSON response: {json_err}"}
        except ValueError as val_err:
             logger.error(f"Invalid JSON structure from LLM: {val_err}. Response text: {response.text}")
             return {"error": f"Invalid JSON structure from LLM: {val_err}"}

    except Exception as e:
        logger.error(f"LLM call failed during entity extraction: {e}", exc_info=True)
        return {"error": f"Error during entity extraction: {e}"}

def split_text_into_chunks(text: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> List[str]:
    """Split a long text into smaller chunks suitable for processing."""
    logger.info(f"Splitting text (length: {len(text)} chars) into chunks (size: {chunk_size}, overlap: {chunk_overlap}).")
    if not text:
        return []
    try:
        splitter = SentenceSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        document = Document(text=text)
        nodes = splitter.get_nodes_from_documents([document])
        chunks = [node.get_content() for node in nodes]
        logger.info(f"Text split into {len(chunks)} chunks.")
        return chunks
    except Exception as e:
        logger.error(f"Error splitting text: {e}", exc_info=True)
        # Fallback to simple splitting if SentenceSplitter fails
        logger.warning("Falling back to simple text splitting.")
        return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size - chunk_overlap)]

def transcribe_audio(audio_file_path: str, language: Optional[str] = None) -> str:
    """Transcribes an audio file using the OpenAI Whisper model.
       Args:
           audio_file_path (str): The path to the audio file (e.g., mp3, wav, m4a).
           language (Optional[str]): The language code (e.g., "en", "es") or full name ("English", "Spanish"). 
                                     If None, Whisper will detect the language.
       Returns:
           str: The transcribed text or an error message.
    """
    logger.info(f"Attempting to transcribe audio file: {audio_file_path}, Language: {language}")
    
    # Check if Whisper is available
    if not WHISPER_AVAILABLE:
        return "Error: openai-whisper library is required but not installed."
        
    # Check if file exists
    if not os.path.exists(audio_file_path):
        logger.error(f"Audio file not found: {audio_file_path}")
        return f"Error: Audio file not found at {audio_file_path}"
        
    # Load the Whisper model (lazy loading)
    model = _load_whisper_model() # Uses default size "base" or WHISPER_MODEL_SIZE env var
    if model is None:
        return "Error: Failed to load Whisper model."

    try:
        # Perform transcription
        # The transcribe function handles various audio formats via ffmpeg
        result = model.transcribe(audio_file_path, language=language)
        transcribed_text = result["text"]
        detected_language = result.get("language", "unknown") # Get detected language if available
        logger.info(f"Audio transcription successful. Detected language: {detected_language}. Text length: {len(transcribed_text)}")
        return transcribed_text
        
    except Exception as e:
        # Check if it might be an ffmpeg issue
        if "ffmpeg" in str(e).lower():
             logger.error(f"Error during transcription, possibly ffmpeg issue: {e}", exc_info=True)
             # Check if ffmpeg is installed using shell command
             try:
                 subprocess.run(["ffmpeg", "-version"], check=True, capture_output=True)
                 # If ffmpeg is installed, the error is likely something else
                 return f"Error during transcription (ffmpeg seems installed): {e}"
             except (FileNotFoundError, subprocess.CalledProcessError):
                 logger.error("ffmpeg command not found or failed. Please ensure ffmpeg is installed and in PATH.")
                 return "Error: ffmpeg not found or not working. Please install ffmpeg."
        else:
             logger.error(f"Unexpected error during transcription: {e}", exc_info=True)
             return f"Error during transcription: {e}"

# --- Tool Definitions ---
summarize_tool = FunctionTool.from_defaults(
    fn=summarize_text,
    name="summarize_text",
    description=(
        "Summarizes a given block of text. Useful for condensing long documents or articles. "
        "Input: text (str), Optional: max_length (int), min_length (int). Output: summary (str) or error."
    ),
)

extract_entities_tool = FunctionTool.from_defaults(
    fn=extract_entities,
    name="extract_entities",
    description=(
        "Extracts named entities (people, organizations, locations, dates, events) from text. "
        "Input: text (str), Optional: entity_types (List[str]). Output: Dict[str, List[str]] or error dict."
    ),
)

split_text_tool = FunctionTool.from_defaults(
    fn=split_text_into_chunks,
    name="split_text_into_chunks",
    description=(
        "Splits a long text document into smaller, overlapping chunks. "
        "Input: text (str), Optional: chunk_size (int), chunk_overlap (int). Output: List[str] of chunks."
    ),
)

# Conditionally create transcribe_audio_tool
transcribe_audio_tool = None
if WHISPER_AVAILABLE:
    transcribe_audio_tool = FunctionTool.from_defaults(
        fn=transcribe_audio,
        name="transcribe_audio_file",
        description=(
            "Transcribes speech from an audio file (e.g., mp3, wav, m4a) into text using Whisper. "
            "Input: audio_file_path (str), Optional: language (str - e.g., \"en\", \"Spanish\"). "
            "Output: transcribed text (str) or error message."
        ),
    )
    logger.info("Audio transcription tool created.")
else:
    logger.warning("Audio transcription tool disabled because openai-whisper is not installed.")

# --- Agent Initialization ---
def initialize_text_analyzer_agent() -> ReActAgent:
    """Initializes the Text Analyzer Agent."""
    logger.info("Initializing TextAnalyzerAgent...")

    # Configuration for the agent's main LLM
    agent_llm_model = os.getenv("TEXT_ANALYZER_AGENT_LLM_MODEL", "models/gemini-1.5-pro")
    gemini_api_key = os.getenv("GEMINI_API_KEY")

    if not gemini_api_key:
        logger.error("GEMINI_API_KEY not found for TextAnalyzerAgent.")
        raise ValueError("GEMINI_API_KEY must be set for TextAnalyzerAgent")

    try:
        llm = GoogleGenAI(api_key=gemini_api_key, model=agent_llm_model)
        logger.info(f"Using agent LLM: {agent_llm_model}")

        # Load system prompt
        default_system_prompt = ("You are TextAnalyzerAgent... [Default prompt content - replace with actual]" # Placeholder
                              )
        system_prompt = load_prompt_from_file("../prompts/text_analyzer_prompt.txt", default_system_prompt)
        if system_prompt == default_system_prompt:
             logger.warning("Using default/fallback system prompt for TextAnalyzerAgent.")

        # Define available tools, including the audio tool if available
        tools = [summarize_tool, extract_entities_tool, split_text_tool]
        if transcribe_audio_tool:
            tools.append(transcribe_audio_tool)
            
        # Update agent description based on available tools
        agent_description = (
            "Analyzes text content. Can summarize text (`summarize_text`), extract named entities (`extract_entities`), "
            "and split long texts (`split_text_into_chunks`)."
        )
        if transcribe_audio_tool:
             agent_description += " Can also transcribe audio files to text (`transcribe_audio_file`)."

        agent = ReActAgent(
            name="text_analyzer_agent",
            description=agent_description,
            tools=tools,
            llm=llm,
            system_prompt=system_prompt,
            can_handoff_to=["planner_agent", "research_agent", "reasoning_agent"], # Example handoffs
        )
        logger.info("TextAnalyzerAgent initialized successfully.")
        return agent

    except Exception as e:
        logger.error(f"Error during TextAnalyzerAgent initialization: {e}", exc_info=True)
        raise

# Example usage (for testing if run directly)
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    logger.info("Running text_analyzer_agent.py directly for testing...")

    # Check required keys
    required_keys = ["GEMINI_API_KEY"]
    missing_keys = [key for key in required_keys if not os.getenv(key)]
    if missing_keys:
        print(f"Error: Required environment variable(s) not set: {', '.join(missing_keys)}. Cannot run test.")
    else:
        try:
            # Test summarization
            print("\nTesting summarization...")
            long_text = """The Industrial Revolution, now also known as the First Industrial Revolution, was a period of global transition of the human economy towards more efficient and stable manufacturing processes that succeeded the Agricultural Revolution, starting from Great Britain, continental Europe and the United States, that occurred during the period from around 1760 to about 1820–1840. This transition included going from hand production methods to machines; new chemical manufacturing and iron production processes; the increasing use of water power and steam power; the development of machine tools; and the rise of the mechanized factory system. The Revolution also saw an unprecedented rise in the rate of population growth."""
            summary = summarize_text(long_text, max_length=50)
            print(f"Summary:\n{summary}")

            # Test entity extraction
            print("\nTesting entity extraction...")
            entities = extract_entities(long_text, entity_types=["EVENT", "GPE", "DATE"])
            print(f"Extracted Entities:\n{entities}")
            
            # Test text splitting
            print("\nTesting text splitting...")
            chunks = split_text_into_chunks(long_text * 3, chunk_size=150, chunk_overlap=30) # Make text longer
            print(f"Split into {len(chunks)} chunks. First chunk:\n{chunks[0]}")
            
            # Test audio transcription (if available)
            if WHISPER_AVAILABLE:
                print("\nTesting audio transcription...")
                # Create a dummy audio file for testing (requires ffmpeg)
                dummy_file = "dummy_audio.mp3"
                try:
                    # Generate a 1-second silent MP3 using ffmpeg
                    subprocess.run(["ffmpeg", "-f", "lavfi", "-i", "anullsrc=r=44100:cl=mono", "-t", "1", "-q:a", "9", "-y", dummy_file], check=True, capture_output=True)
                    print(f"Created dummy audio file: {dummy_file}")
                    transcript = transcribe_audio(dummy_file)
                    print(f"Transcription Result: '{transcript}' (Expected: empty or silence markers)")
                    os.remove(dummy_file) # Clean up dummy file
                except Exception as ffmpeg_err:
                    print(f"Could not create/test dummy audio file (ffmpeg required): {ffmpeg_err}")
            else:
                print("\nSkipping audio transcription test as openai-whisper is not available.")

            # Initialize the agent (optional)
            # test_agent = initialize_text_analyzer_agent()
            # print("\nText Analyzer Agent initialized successfully for testing.")

        except Exception as e:
            print(f"Error during testing: {e}")