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 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


# 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

# --- 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="gemini-2.5-pro-preview-03-25", temperature=0.05)
        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)]

# --- 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."
    ),
)

# --- 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", "gemini-2.5-pro-preview-03-25")
    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="gemini-2.5-pro-preview-03-25", temperature=0.05)
        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]
            
        # 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`)."
        )

        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", "verifier_agent", "advanced_validation_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]}")

            # 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}")