import os
import math

import smolagents
import smolagents.tools


# Rectangular section properties tool
class RectSectionPropsTool(smolagents.tools.Tool):
    name = "rect_section_props"
    description = "Rectangular section properties."
    inputs = {
        "b": {"type": "number", "description": "Width in meters"},
        "h": {"type": "number", "description": "Height in meters"},
    }
    output_type = "object"

    def forward(self, b: float, h: float) -> dict:
        if b <= 0 or h <= 0:
            raise ValueError("b and h must be positive.")
        A = b * h
        I = b * (h**3) / 12.0
        S = I / (h / 2.0)  # = b*h^2/6
        return {"A": A, "I": I, "S": S}


# Circle section properties tool
class CircleSectionPropsTool(smolagents.tools.Tool):
    name = "circle_section_props"
    description = "Circular section properties."
    inputs = {
        "d": {"type": "number", "description": "Diameter in meters"},
    }
    output_type = "object"

    def forward(self, d: float) -> dict:
        if d <= 0:
            raise ValueError("d must be positive.")
        A = math.pi * (d**2) / 4.0
        I = math.pi * (d**4) / 64.0
        S = I / (d / 2.0)  # = π d^3 / 32
        return {"A": A, "I": I, "S": S}


# Mass tool
class BeamMassTool(smolagents.tools.Tool):
    name = "beam_mass"
    description = "Beam mass [kg]."
    inputs = {
        "L": {"type": "number", "description": "Length in meters"},
        "rho": {"type": "number", "description": "Material density in kg/m^3"},
        "A": {
            "type": "number",
            "description": "Cross-sectional area in m^2 (optional)",
        },
    }
    output_type = "number"

    def forward(self, L: float, rho: float, A: float) -> float:
        if L <= 0 or rho <= 0 or A <= 0:
            raise ValueError("L, rho, and A must be positive.")
        return rho * A * L


# Deflection tool
class BeamDeflectionTool(smolagents.tools.Tool):
    name = "beam_deflection"
    description = "Max deflection δ [m] for simple cases."
    inputs = {
        "P": {"type": "number", "description": "Point load in Newtons"},
        "L": {"type": "number", "description": "Span length in meters"},
        "E": {"type": "number", "description": "Young's modulus in Pascals"},
        "I": {"type": "number", "description": "Second moment in m^4 (optional)"},
    }
    output_type = "number"

    def forward(self, P: float, L: float, E: float, I: float) -> float:
        if any(x <= 0 for x in [P, L, E, I]):
            raise ValueError("P, L, E, I must be positive.")
        return P * (L**3) / (3.0 * E * I)


# Bending stress tool
class BeamBendingStressTool(smolagents.tools.Tool):
    name = "beam_bending_stress"
    description = "Max bending stress σ_max [Pa]."
    inputs = {
        "M": {"type": "number", "description": "Bending moment in N·m"},
        "I": {"type": "number", "description": "Second moment in m^4 (optional)"},
        "c": {"type": "number", "description": "Neutral axis distance in m (optional)"},
    }
    output_type = "number"

    def forward(self, M: float, I: float, c: float) -> float:
        if M <= 0:
            raise ValueError("M must be positive.")
        return M * c / I


def generate_beam_agent() -> smolagents.CodeAgent:
    # Define the agent with all of these tools.
    return smolagents.CodeAgent(
        tools=[
            BeamMassTool(),
            BeamDeflectionTool(),
            BeamBendingStressTool(),
            RectSectionPropsTool(),
            CircleSectionPropsTool(),
        ],
        model=smolagents.models.OpenAIServerModel(
            model_id=os.getenv("AGENT_MODEL", ""),
            api_base=os.getenv("UPSTREAM_OPENAI_BASE", "").rstrip("/"),
            api_key=os.getenv("OPENAI_API_KEY"),
        ),
        instructions="You are a beam design agent with custom tools to support that task. Use those tools whenever feasible.",
        name="beam_design_agent",
        description="Agent specialized in beam design calculations.",
        add_base_tools=False,
        max_steps=12,
    )