import os

# 1) Set up environment so HF caches models in /tmp/hf
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf"
os.environ["HF_HOME"] = "/tmp/hf"

from langchain_community.llms import HuggingFacePipeline
from langchain.agents import initialize_agent
from langchain.agents.agent_types import AgentType
from langchain.tools import BaseTool
from transformers import pipeline

# 2) Use a smaller model + force_download=True to avoid partial downloads
pipe = pipeline(
    "text-generation",
    model="bigscience/bloom-560m",
    tokenizer="bigscience/bloom-560m",
    device=-1,               # CPU only
    max_new_tokens=256,      # A smaller limit speeds up generation
    force_download=True      # Forces a fresh download if needed
)

# 3) Wrap the pipeline in a LangChain LLM
llm = HuggingFacePipeline(pipeline=pipe)

# 4) Define your custom tool for generating the machinery report
class MachineryReportTool(BaseTool):
    name = "machinery_report"
    description = (
        "This tool compiles a detailed report on the mini construction equipment project, "
        "including specifications, market analysis, and production considerations."
    )

    def _run(self, query: str) -> str:
        report = """
        ## Comprehensive Report on Mini Construction Equipment Project
        ### Overview
        The project involves designing and building construction machinery tailored for the local South African market. 
        The focus is on developing cost-effective, high-performance machines that can compete with expensive American-made 
        equipment while leveraging local manufacturing strengths.
        ### 1. Equipment Details
        - **Basic Gas-Powered Auction Unit**
          - Price: $3,700
          - Engine: 14 horsepower gas motor
          - Configuration: Two pump system
          - Capabilities: Digging, Scooping, Auger Operation, Trenching, Pallet Fork Operation
        - **DRT 450**
          - Original Price: $8,500 (current base price: $5,500-6,300)
          - Engine: Twin cylinder Honda gas motor
          - Pump System: Triple pump configuration
          - Compatible Attachments: Mulcher, Brush Cutter, Harley Rake, Trencher
        - **Mini Skid Steer (New Acquisition)**
          - Price: $15,000
          - Engine: Kubota diesel
          - Control: Pilot controls with vertical lift arms
          - Notable Feature: Mulcher attachment ($2,000)
        - **Mini Excavators (Kimron Units)**
          - Unit 1: One-ton unit with Briggs & Stratton engine
          - Unit 2: 3,500 lb unit with Yanmar diesel engine
          - Supplier: K&R Equipment (Oklahoma)
        ### 2. Market Comparison and Cost Impact
        - **American Equivalents:** Bobcat MT100, Kubota SCCL1000, Ditch Witch SK900 ($38,000 to $45,000+).
        - **Tariff Impact:** 
          - DRT 450: ~$6,300 base + 25% + additional 10% (~$630) = ~$9,000.
          - Mini Skid Steer: ~$9,300 + ~$5,000 shipping + 10% tariff (~$930) = ~$16,000.
        - **Competitive Advantage:** 
          Even with tariffs, Chinese equipment can be 25-30% the cost of American machines.
        ### 3. Designing and Building Locally in South Africa
        - **Local Manufacturing Advantages:**
          - "Made in South Africa" branding for local pride.
          - Custom designs for local farming and construction.
        - **Production Considerations:**
          - Assemble components locally to reduce tariffs.
          - Partner with local engineering firms.
          - Comply with SA standards (and international if exporting).
        - **Market Entry Strategies:**
          - Pilot production, gather feedback, scale.
          - Target local contractors/farmers.
          - Check government incentives for local manufacturing.
        ### 4. Next Steps & Recommendations
        - **R&D:** Prototypes + real-world testing.
        - **Supply Chain:** Local suppliers, modular designs.
        - **Financing:** Bank loans, grants, investor funding.
        - **Compliance:** Prepare for certifications and warranties.
        ### Conclusion
        Building "Made in South Africa" machinery meets a real need for affordable, durable equipment, 
        potentially disrupting a market dominated by high-priced American brands.
        """
        return report

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

tools = [MachineryReportTool()]

# 5) Initialize the agent with max_iterations=1 to avoid looping
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    max_iterations=1
)

# 6) Run the agent with your query
query = (
    "Using the available tools, please compile a detailed research report on the "
    "mini construction equipment project, including design, market analysis, and production "
    "specifications based on the provided machinery details."
)

result = agent.run(query)
print("\n===== AGENT OUTPUT =====")
print(result)