from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI

# Fetch mello roos taxes for addresses in Sacramento County
@tool
def sacramento_county_mello_roos_tool(street_number:str, street_name:str, city:str)-> str: #it's import to specify the return type
    #Keep this format for the description / args / args description but feel free to modify the tool
    """A tool that fetches the tax data for addresses in the Sacramento county and returns the levy portions of the property taxes including the Mello Roos taxes 
    Args:
        street_number: Street number of the address
        street_name: Street name of the address
        city: City of the address
    """
    parcel_api_url = "https://services1.arcgis.com/5NARefyPVtAeuJPU/arcgis/rest/services/Parcels/FeatureServer/0/query"
    where_clause = f"STREET_NBR='{street_number}' AND STREET_NAM='{street_name}' AND CITY='{city}'"
    params = {
        'where': where_clause,
        'outFields': '*',
        'f': 'json'
    }

    # Fetch parcel data
    parcel_response = requests.get(parcel_api_url, params=params)
    if parcel_response.status_code != 200:
        return jsonify({"error": f"Parcel API request failed with status code {parcel_response.status_code}, response: {parcel_response.text}"}), 400

    parcel_data = parcel_response.json()

    # Ensure parcel data is valid
    if 'features' not in parcel_data or len(parcel_data['features']) == 0:
        return jsonify({"error": "No parcel data found"}), 400

    parcel_number = parcel_data['features'][0]['attributes'].get('APN')
    if not parcel_number:
        return jsonify({"error": "Parcel number not found"}), 400

    # Fetch bill details using the parcel number
    bill_summary_url = f"https://eproptax.saccounty.net/servicev2/eproptax.svc/rest/BillSummary?parcel={parcel_number}"
    bill_response = requests.get(bill_summary_url)
    if bill_response.status_code != 200:
        return jsonify({"error": f"Bill API request failed with status code {bill_response.status_code}, response: {bill_response.text}"}), 400

    bill_data = bill_response.json()

    # Ensure bill data is valid
    if not bill_data.get('Success'):
        return jsonify({"error": "Failed to retrieve bill data"}), 400

    # Get the bill number and roll date
    bill_number = bill_data['Bills'][0]['BillNumber']
    roll_date = bill_data['Bills'][0]['RollDate']

    # Fetch levy details using the bill number and current year
    levy_url = f"https://eproptax.saccounty.net/servicev2/eproptax.svc/rest/DirectLevy?rollYear={roll_date}&billNumber={bill_number}"
    levy_response = requests.get(levy_url)
    if levy_response.status_code != 200:
        return jsonify({"error": f"Levy API request failed with status code {levy_response.status_code}, response: {levy_response.text}"}), 400

    levy_data = levy_response.json()

    # Ensure levy data is valid
    if not levy_data.get('Success'):
        return jsonify({"error": "Failed to retrieve levy data"}), 400

    # Prepare the response
    response_data = {
        "total_amount": levy_data['BillAmount'],
        "levy_total": levy_data['LevyTotal'],
        "levies": levy_data['Levies']
    }

    return str(response_data)

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[sacramento_county_mello_roos_tool, image_generation_tool, final_answer], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()