from fastapi import FastAPI,Header
import uvicorn
"""
Simple LangGraph Agent using @tool decorator

Clean implementation with decorator-based tool definitions.
"""

from typing import Annotated, TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.tools import tool
from langchain_aws import ChatBedrock
app = FastAPI()

# Define tools using @tool decorator
rules="""- Mulheres acima de 45 anos ou menopausada
- Homens com mais de 70 anos;
- Osteogênese imperfeita (para esta patologia, poderá haver a liberação de (02) dois 
exames ao ano - cada 180 dias);
- RX com osteopenia ou fratura patológica;
- Antecedente pessoal de fratura após os 40 anos: punho, ombros, vértebras, quadril;
- Parente de primeiro grau com osteoporose.
- Mulheres com massa corporal <20kg/m2 ou peso < 57,8kg;
- Menopausa antes dos 45 anos ou hipogonasismo crônico (falência ovariana 
precoce);
- Uso de glicocorticóides (>=7,5 prednizona/ dia equivalente por mais três meses, ou 
presença de síndrome de cushing;
- Hiperparatireoidismo primário;
- Uso prolongado de anticonvulsivantes (< 10 anos);
- Síndrome de má absorção crônica ou desnutrição doenças inflamatória intestinal 
(independente da causa: bariatrica, celiacos, intolerancia a lactose).
- Quimioterapia, se sobrevida esperada for longa (< 5 anos);
- Diminuição documentada de altura;
- Presença de cifose após menopausa.
- Imobilização prolongada"""
SYSTEM_PROMPT="""
You are a assisant, your job is to aprove or not a procedure based on the following rules:"""
+rules+"""
Your input will be a json, evaluate it based on the rules and return:
Aproved: If one of the criteira is met
Reproved: If not a single one of the criterias are met."""
@tool
def add(a: int, b: int) -> int:
    """Add two numbers together.
    
    Args:
        a: First number
        b: Second number
    """
    return a + b


@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers.
    
    Args:
        a: First number
        b: Second number
    """
    return a * b


@tool
def get_word_length(word: str) -> int:
    """Get the length of a word.
    
    Args:
        word: The word to measure
    """
    return len(word)


@tool
def search_info(topic: str) -> str:
    """Search for information about a topic (mock implementation).
    
    Args:
        topic: The topic to search for
    """
    # Mock response - replace with actual search/API
    return f"Information about {topic}: This is a mock response. In production, this would return real data."


# Define agent state
class AgentState(TypedDict):
    messages: Annotated[list, add_messages]


# Define tools list
tools = [add, multiply, get_word_length, search_info]


# Agent node
def call_model(state: AgentState):
    """Call the LLM with current state and tools."""
    
    model = ChatBedrock(
        model_id="arn:aws:bedrock:us-east-2:232048051668:application-inference-profile/uy4xskop19zn",
        region_name="us-east-2",
        provider="anthropic"
    )
    
    model_with_tools = model.bind_tools(tools)
    
    messages = [
        SystemMessage(content=SYSTEM_PROMPT)
    ] + state["messages"]
    
    response = model_with_tools.invoke(messages)
    
    return {"messages": [response]}


# Build the graph
def create_agent():
    """Create and compile the agent graph."""
    
    workflow = StateGraph(AgentState)
    
    # Add nodes
    workflow.add_node("agent", call_model)
    workflow.add_node("tools", ToolNode(tools))
    
    # Add edges
    workflow.add_edge(START, "agent")
    workflow.add_conditional_edges("agent", tools_condition)
    workflow.add_edge("tools", "agent")
    
    return workflow.compile()


# Main execution


@app.post("/")
async def root(json:str= Header(...)):
    agent = create_agent()
    query=json
    result = agent.invoke(
            {"messages": [HumanMessage(content=query)]},
            config={"recursion_limit": 10}
        )
        
    final_message = result["messages"][-1]
    return {"status": "success", "message": final_message.content}

@app.get("/health")
async def health():
    return {"status": "healthy"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)