from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import pytz
import yaml
from skyfield.api import load, Topos, load_file
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Define Zodiac signs and their boundaries (0° to 360° ecliptic longitude)
ZODIAC_SIGNS = [
    ("Aries", 0, 30),
    ("Taurus", 30, 60),
    ("Gemini", 60, 90),
    ("Cancer", 90, 120),
    ("Leo", 120, 150),
    ("Virgo", 150, 180),
    ("Libra", 180, 210),
    ("Scorpio", 210, 240),
    ("Sagittarius", 240, 270),
    ("Capricorn", 270, 300),
    ("Aquarius", 300, 330),
    ("Pisces", 330, 360),
]

# Moon phase boundaries (0° to 360° phase angle) for display purposes
MOON_PHASES = [
    ("New Moon", 0, 45),
    ("Waxing Crescent", 45, 90),
    ("First Quarter", 90, 135),
    ("Waxing Gibbous", 135, 180),
    ("Full Moon", 180, 225),
    ("Waning Gibbous", 225, 270),
    ("Last Quarter", 270, 315),
    ("Waning Crescent", 315, 360),
]

# Fertility sign coefficients (applicable to all plants)
FERTILITY_SIGN_COEFFS = {
    "Aries": 1,
    "Taurus": 2,
    "Gemini": 0,
    "Cancer": 2,
    "Leo": 1,
    "Virgo": 0,
    "Libra": 0.5,
    "Scorpio": 1.5,
    "Sagittarius": 1,
    "Capricorn": 1,
    "Aquarius": 0,
    "Pisces": 2,
}

# Pruning sign coefficients (applicable to all plants)
PRUNING_SIGN_COEFFS = {
    "Aries": 1,
    "Taurus": 0,
    "Gemini": 2,
    "Cancer": 0,
    "Leo": 1,
    "Virgo": 2,
    "Libra": 1.5,
    "Scorpio": 0.5,
    "Sagittarius": 1,
    "Capricorn": 1,
    "Aquarius": 2,
    "Pisces": 0,
}

# Fertility phase coefficients for above-ground plants
FERTILITY_PHASE_COEFFS_ABOVE = {
    "New Moon": 0,
    "Waxing Moon": 1,
    "Full Moon": 0,
    "Waning Moon": 0.5,
}

# Fertility phase coefficients for root crops
FERTILITY_PHASE_COEFFS_ROOT = {
    "New Moon": 0,
    "Waxing Moon": 0.5,
    "Full Moon": 0,
    "Waning Moon": 1,
}

# Pruning phase coefficients
PRUNING_PHASE_COEFFS = {
    "New Moon": 0,
    "Waxing Moon": 1,
    "Full Moon": 0,
    "Waning Moon": 0.5,
}

@tool
def get_moon_info(date_time: str) -> dict:
    """
    Returns Moon's Zodiac position, phase, and fertility and pruning indices for the given date/time.

    Args:
        date_time (str): ISO 8601 formatted datetime (YYYY-MM-DDTHH:MM:SS)

    Returns:
        dict: {
            "zodiac_position": "Leo 15°30'",
            "moon_phase": "Waxing Gibbous",
            "fertility_above_ground": 2.0,
            "fertility_root_crop": 1.5,
            "pruning": 2.0
        }
    where fertility indices ranges from 0 - minimal expected fertility, to 3.0 - most favorable fertility for platining, 
    and depends on type of plant (root crop or produce above ground);
    pruning indices ranges from 0 - when pruning causes damages to the tree, 2.0 - only minimum or sanitary pruning allowed,
    to 3.0 - most favorable time for pruning;
    """
    try:
        # Parse input datetime and localize to UTC
        user_time = datetime.datetime.strptime(date_time, "%Y-%m-%dT%H:%M:%S")
        user_time = pytz.utc.localize(user_time)

        # Load ephemeris and timescale
        planets = load('https://naif.jpl.nasa.gov/pub/naif/generic_kernels/spk/planets/a_old_versions/de421.bsp')
        ts = load.timescale()
        t = ts.from_datetime(user_time)

        # Define celestial bodies
        earth = planets['earth']
        moon = planets['moon']
        sun = planets['sun']

        # Calculate Moon's ecliptic longitude
        astrometric = earth.at(t).observe(moon)
        _, _, ecliptic_lon = astrometric.ecliptic_latlon()
        lon_deg = ecliptic_lon.degrees % 360

        # Calculate phase angle (angle between Moon and Sun from Earth)
        phase_angle = earth.at(t).observe(moon).apparent().phase_angle(sun).degrees

        # Determine Zodiac sign and position
        zodiac_sign = "Unknown"
        position_degrees = 0
        for sign, start, end in ZODIAC_SIGNS:
            if start <= lon_deg < end:
                zodiac_sign = sign
                position_degrees = lon_deg - start
                break

        # Format position to degrees and minutes
        degrees = int(position_degrees)
        minutes = int((position_degrees % 1) * 60)
        position_str = f"{zodiac_sign} {degrees}°{minutes:02}'"

        # Determine moon phase for display
        moon_phase = "Unknown"
        for phase, start, end in MOON_PHASES:
            if start <= phase_angle < end:
                moon_phase = phase
                break

        # Determine phase category for indices with 15° orbis for New and Full Moon
        if (phase_angle >= 345 or phase_angle < 15):
            phase_category = "New Moon"  # 345° to 15° (30° total orbis)
        elif 15 <= phase_angle < 165:
            phase_category = "Waxing Moon"
        elif 165 <= phase_angle < 195:
            phase_category = "Full Moon"  # 165° to 195° (30° total orbis)
        elif 195 <= phase_angle < 345:
            phase_category = "Waning Moon"
        else:
            phase_category = "Unknown"

        # Calculate fertility and pruning indices
        if zodiac_sign in FERTILITY_SIGN_COEFFS and phase_category in FERTILITY_PHASE_COEFFS_ABOVE:
            fertility_above_ground = FERTILITY_SIGN_COEFFS[zodiac_sign] + FERTILITY_PHASE_COEFFS_ABOVE[phase_category]
            fertility_root_crop = FERTILITY_SIGN_COEFFS[zodiac_sign] + FERTILITY_PHASE_COEFFS_ROOT[phase_category]
            pruning = PRUNING_SIGN_COEFFS[zodiac_sign] + PRUNING_PHASE_COEFFS[phase_category]
        else:
            fertility_above_ground = None
            fertility_root_crop = None
            pruning = None

        return {
            "zodiac_position": position_str,
            "moon_phase": moon_phase,
            "fertility_above_ground": fertility_above_ground,
            "fertility_root_crop": fertility_root_crop,
            "pruning": pruning
        }

    except Exception as e:
        return {"error": f"Error: {str(e)}"}

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Returns the current local time in the specified timezone with description.

    Args:
        timezone (str): A string representing a valid timezone (e.g., 'UTC')

    Returns:
        str: Formatted local time with timezone description
    """
    try:
        tz = pytz.timezone(timezone)
        now = datetime.datetime.now(tz)
        return f"Local time in {timezone}: {now.strftime('%Y-%m-%d %H:%M:%S')}"
    except Exception as e:
        return f"Error: {str(e)}"

@tool
def get_current_time_raw(timezone: str) -> str:
    """
    Returns current local time in specified timezone as ISO 8601 string.

    Args:
        timezone (str): A string representing a valid timezone (e.g., 'UTC')

    Returns:
        str: Datetime in ISO 8601 format (YYYY-MM-DDTHH:MM:SS)
    """
    try:
        tz = pytz.timezone(timezone)
        now = datetime.datetime.now(tz)
        return now.strftime("%Y-%m-%dT%H:%M:%S")
    except Exception as e:
        return f"Error: {str(e)}"

# Model configuration
final_answer = FinalAnswerTool()
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/",
    custom_role_conversions=None,
)

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

# Load prompt templates
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# Initialize agent with all tools
agent = CodeAgent(
    model=model,
    tools=[final_answer, get_moon_info, get_current_time_in_timezone, get_current_time_raw, image_generation_tool],
    max_steps=6,
    verbosity_level=1,
    prompt_templates=prompt_templates
)

if __name__ == "__main__":
    GradioUI(agent).launch()