import os
import json
from dotenv import load_dotenv
from arctic_orchestra.Agents.base import Agent
from arctic_orchestra.Agents.sequential_agent import SequentialAgent
from arctic_orchestra.Agents.loop_agent import LoopAgent
from arctic_orchestra.Agents.router_agent import RoutingAgent
# Load env
load_dotenv()

# ---------------------------------------------------------------------
# TOOL FUNCTIONS FOR TRIP PLANNER
# ---------------------------------------------------------------------

def get_weather(city: str) -> str:
    """Returns fake weather based on city for travel planning."""
    data = {
        "Paris": "Rainy, 12°C",
        "London": "Cloudy, 10°C",
        "Tokyo": "Sunny, 18°C",
        "Delhi": "Clear, 22°C",
        "Mumbai": "Humid, 28°C",
    }
    result = data.get(city, "Weather data unavailable")
    return json.dumps({"city": city, "weather": result})


def get_distance(city1: str, city2: str) -> str:
    """Returns approximate distance in km between two cities."""
    distances = {
        ("Paris", "London"): 344,
        ("London", "Paris"): 344,
        ("Tokyo", "Delhi"): 5820,
        ("Delhi", "Tokyo"): 5820,
        ("Mumbai", "Delhi"): 1148,
        ("Delhi", "Mumbai"): 1148,
    }
    km = distances.get((city1, city2), 1000)
    mode = "flight" if km > 800 else "train"
    return json.dumps({"city1": city1, "city2": city2, "km": km, "recommended_mode": mode})


def estimate_budget(days: int, distance_km: int) -> str:
    """Estimate total trip budget based on days and travel distance."""
    try:
        hotel_cost = days * 120
        travel_cost = 0.1 * distance_km
        total = hotel_cost + travel_cost
        return json.dumps({"days": days, "distance": distance_km, "total_budget_usd": total})
    except Exception as e:
        return json.dumps({"error": str(e)})


# ---------------------------------------------------------------------
# INSTANTIATE AGENTS (using your new Agent class)
# ---------------------------------------------------------------------

def make_agent(name, identity, instruction, tools=None):
    return Agent(
        model_name= "openrouter/qwen/qwen3-30b-a3b",
        api_key=os.getenv("OPENROUTER_API_KEY"),
        name=name,
        identity=identity,
        instruction=instruction,
        tools=tools or {},
        debug=True
    )


# 1. Extracts trip request into structured info
LocationAgent = make_agent(
    "LocationAgent",
    "You extract travel details.",
    "Identify origin, destination, number of days, and user preferences.",
)

# 2. Calls weather tool
WeatherAgent = make_agent(
    "WeatherAgent",
    "You provide weather information.",
    "Use the get_weather tool when weather is required.",
    tools={"get_weather": get_weather}
)

# 3. Calls distance tool
DistanceAgent = make_agent(
    "DistanceAgent",
    "You compute travel distances.",
    "Always use get_distance tool to compute distance.",
    tools={"get_distance": get_distance}
)

# 4. Calls budget tool
BudgetAgent = make_agent(
    "BudgetAgent",
    "You estimate travel budget.",
    "Always use estimate_budget tool when cost calculation is needed.",
    tools={"estimate_budget": estimate_budget}
)

# 5. Final synthesis
FinalPlannerAgent = make_agent(
    "FinalPlannerAgent",
    "You create a final trip plan.",
    "Summarize full results into a clean travel itinerary." \
    f"Plan from Location agent {LocationAgent.final_response} \n"
    f"Plan from Weather Agent {WeatherAgent.final_response} \n"
    f"Plan from Distance Agent {DistanceAgent.final_response} \n"
    f"Plan from Budget Agent {BudgetAgent.final_response} \n"

)


# ---------------------------------------------------------------------
# ORCHESTRATOR
# ---------------------------------------------------------------------

TripPlanner = LoopAgent(
    name="TripPlannerLoop",
    agents=[LocationAgent, WeatherAgent, DistanceAgent, BudgetAgent, FinalPlannerAgent],
    max_iterations=3,
    window_size=2,
    max_context_chars=8000,
    exit_allowed_agents=[FinalPlannerAgent],
    exit_instructions="use the exit_tool only if plans from all the agents are present to make the summary . if they are missing dont use the tool exit_loop."
)

TripPlanner = RoutingAgent(
        model_name= "openrouter/x-ai/grok-4.1-fast",
        api_key=os.getenv("OPENROUTER_API_KEY"),
        agents=[WeatherAgent, DistanceAgent, BudgetAgent, FinalPlannerAgent],
        additional_routing_instructions="""CRITICAL EXECUTION RULES:
    1. You must execute the tools in the exact order listed above (Weather -> Distance -> Budget -> FinalPlannerAgent).
    2. Collect the output from every single agent.
    3. call only one tool at a time.
    4. Do NOT call the 'finalplanneragent' until you have successful outputs from all previous agents.
    5. Pass the combined findings from Location, Weather, Distance, and Budget into the 'finalplanneragent' so it can generate the final summary.
    6. donot ouput anything just continuously call teh tools once finished calling in teh order then ouput
    """,
        debug=True
).build_router_agent()




# ---------------------------------------------------------------------
# RUN POC TEST
# ---------------------------------------------------------------------

user_query = """
Plan a 5-day trip from Delhi to Tokyo.
I care about weather, recommended travel mode, and total budget.
"""

print("\n\n=== RUNNING TRIP PLANNER ORCHESTRATION ===\n")
result = TripPlanner.run(user_query)

print("\n\n=== FINAL TRIP PLAN ===")
print(result)
