import os
import sys
import logging
from typing import Optional, Any
from typing_extensions import override

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

def setup_environment():
    """Setup environment variables and cache directories"""
    try:
        cache_dir = "/tmp/hf"
        os.environ["TRANSFORMERS_CACHE"] = cache_dir
        os.environ["HF_HOME"] = cache_dir
        
        # Create cache directory if it doesn't exist
        os.makedirs(cache_dir, exist_ok=True)
        logger.info(f"Cache directory set up at {cache_dir}")
        return True
    except Exception as e:
        logger.error(f"Failed to setup environment: {str(e)}")
        return False

def initialize_model():
    """Initialize the model and pipeline with error handling"""
    try:
        from transformers import pipeline
        logger.info("Loading model pipeline...")
        
        pipe = pipeline(
            "text-generation",
            model="bigscience/bloom-560m",
            tokenizer="bigscience/bloom-560m",
            device=-1,
            max_new_tokens=256,
            force_download=True,
            low_cpu_mem_usage=True,
        )
        logger.info("Model pipeline loaded successfully")
        return pipe
    except Exception as e:
        logger.error(f"Failed to initialize model: {str(e)}")
        return None

def setup_agent(pipe):
    """Setup LangChain agent with the model pipeline"""
    try:
        from langchain_huggingface import HuggingFacePipeline  # Updated import
        from langchain.agents import initialize_agent
        from langchain.agents.agent_types import AgentType
        from langchain.tools import BaseTool

        # Initialize LLM
        llm = HuggingFacePipeline(pipeline=pipe)

        # Define the tool with proper type annotations
        class MachineryReportTool(BaseTool):
            name: str = "machinery_report"  # Type annotation added
            description: str = (  # Type annotation added
                "Generates a report on mini construction equipment including "
                "specifications and market analysis."
            )

            @override
            def _run(self, query: str) -> str:
                return """
                Mini Construction Equipment Report:
                1. Basic Gas-Powered Unit: $3,700, 14HP
                2. DRT 450: $5,500-6,300, Honda motor
                3. Mini Skid Steer: $15,000, Kubota diesel
                Market Analysis: 30-40% cheaper than US equivalents
                """

            @override
            def _arun(self, query: str) -> Any:
                raise NotImplementedError("Async not supported.")

        # Initialize agent with minimal settings
        agent = initialize_agent(
            tools=[MachineryReportTool()],
            llm=llm,
            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            verbose=True,
            max_iterations=1,
            early_stopping_method="generate"
        )
        logger.info("Agent initialized successfully")
        return agent
    except Exception as e:
        logger.error(f"Failed to setup agent: {str(e)}")
        return None

def run_agent_query(agent, query: str) -> Optional[str]:
    """Run a query through the agent with error handling"""
    try:
        logger.info("Running agent query...")
        result = agent.run(query)
        logger.info("Query completed successfully")
        return result
    except Exception as e:
        logger.error(f"Error running agent query: {str(e)}")
        return f"Error processing query: {str(e)}"

def main():
    """Main function to run the agent"""
    logger.info("Starting application...")
    
    # Setup environment
    if not setup_environment():
        return
    
    # Initialize model
    pipe = initialize_model()
    if pipe is None:
        return
    
    # Setup agent
    agent = setup_agent(pipe)
    if agent is None:
        return
    
    # Define query
    query = (
        "Please provide a brief overview of the mini construction equipment "
        "project, focusing on key specifications and market positioning."
    )
    
    # Run query
    result = run_agent_query(agent, query)
    
    # Print result
    if result:
        print("\n===== AGENT OUTPUT =====")
        print(result)
        print("=======================")

if __name__ == "__main__":
    main()