from smolagents import CodeAgent, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import os
from dotenv import load_dotenv
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Load environment variables
load_dotenv()

# API Configuration
THERUNDOWN_API_KEY = os.getenv('THERUNDOWN_API_KEY')
if not THERUNDOWN_API_KEY:
    raise ValueError("THERUNDOWN_API_KEY not found in environment variables")

@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 valid timezone string (e.g., 'America/New_York').
    """
    try:
        tz = pytz.timezone(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)}"

@tool
def get_nba_matches(target_date: str = None) -> str:
    """
    A tool that retrieves upcoming NBA matches using TheRundown API.
    Uses the /openers endpoint for the specified date.
    
    Args:
        target_date: Optional date string in YYYY-MM-DD format. If not provided, uses today's date.
    
    Returns:
        A human-readable string listing the upcoming NBA matches with detailed odds information.
        For special events like the All-Star Game, includes additional event details.
    """
    import requests
    import datetime
    from typing import Dict, Any
    import json

    # Use provided date or today's date
    if target_date is None:
        target_date = datetime.date.today().strftime("%Y-%m-%d")
    
    # API configuration
    url = f"https://api.apilayer.com/therundown/sports/4/openers/{target_date}"
    params = {
        "offset": "0",
        "include": "scores"
    }
    headers = {
        "apikey": THERUNDOWN_API_KEY
    }
    
    try:
        response = requests.get(url, headers=headers, params=params)
        response.raise_for_status()
        
        data = response.json()
        if not data.get("events"):
            return f"No NBA matches found for {target_date}."
            
        matches = []
        is_special_event = False
        event_type = "Regular Season"
        
        for event in data["events"]:
            # Extract event date and time
            event_date = event.get("event_date", "N/A")
            if event_date != "N/A":
                try:
                    dt = datetime.datetime.strptime(event_date, "%Y-%m-%dT%H:%M:%SZ")
                    event_date = dt.strftime("%Y-%m-%d %H:%M UTC")
                except ValueError:
                    pass
            
            # Get schedule information
            schedule = event.get("schedule", {})
            event_headline = schedule.get("event_headline", "")
            
            # Check if this is a special event
            if "All-Star" in event_headline:
                is_special_event = True
                event_type = "NBA All-Star Game"
            
            # Extract basic game info
            teams = event.get("teams_normalized", [])
            if not teams:
                teams = event.get("teams", [])
            
            if len(teams) < 2:
                continue
                
            away_team = teams[0] if teams[0].get("is_away") else teams[1]
            home_team = teams[1] if teams[0].get("is_away") else teams[0]
            
            # Get the lines info (odds)
            lines = event.get("lines", {})
            
            # Try different bookmakers in order of preference
            bookmaker = None
            for bookmaker_id in ["3", "2", "12", "22", "23"]:
                if bookmaker_id in lines:
                    bookmaker = lines[bookmaker_id]
                    break
            
            if not bookmaker:
                bookmaker = next(iter(lines.values())) if lines else {}
            
            # Extract odds information
            spread_info = bookmaker.get("spread", {})
            total_info = bookmaker.get("total", {})
            moneyline_info = bookmaker.get("moneyline", {})
            
            # Get venue and broadcast information
            score_info = event.get("score", {})
            venue_name = score_info.get("venue_name", "N/A")
            venue_location = score_info.get("venue_location", "N/A")
            broadcast = score_info.get("broadcast", "N/A")
            
            # Format game information
            game_info = []
            
            # Add event type header if it's a special event
            if is_special_event and event_headline:
                game_info.extend([
                    "🏀 " + "=" * 30 + " NBA ALL-STAR GAME " + "=" * 30 + " 🏀",
                    f"Event: {event_headline}",
                    f"Date: {event_date}",
                ])
            else:
                game_info.extend([
                    "🏀 " + "=" * 35 + " NBA GAME " + "=" * 35 + " 🏀",
                    f"Date: {event_date}",
                ])
            
            # Add team information
            game_info.append(
                f"Teams: {away_team.get('name', 'Unknown')} vs {home_team.get('name', 'Unknown')}"
            )
            
            # Add venue and broadcast
            game_info.extend([
                f"Venue: {venue_name} ({venue_location})",
                f"Broadcast: {broadcast}",
            ])
            
            # Add betting information if available
            if any(v != 0.0001 for v in [
                moneyline_info.get("moneyline_away", 0),
                moneyline_info.get("moneyline_home", 0),
                spread_info.get("point_spread_away", 0),
                total_info.get("total_over", 0)
            ]):
                game_info.append("\nBetting Lines:")
                
                if moneyline_info.get("moneyline_away", 0) != 0.0001:
                    game_info.append(
                        f"Moneyline: {away_team.get('name')} ({moneyline_info.get('moneyline_away', 'N/A')}) | "
                        f"{home_team.get('name')} ({moneyline_info.get('moneyline_home', 'N/A')})"
                    )
                
                if spread_info.get("point_spread_away", 0) != 0.0001:
                    game_info.append(
                        f"Spread: {away_team.get('name')} {spread_info.get('point_spread_away', 'N/A')} "
                        f"({spread_info.get('point_spread_away_money', 'N/A')})"
                    )
                
                if total_info.get("total_over", 0) != 0.0001:
                    game_info.append(
                        f"Total: O/U {total_info.get('total_over', 'N/A')} "
                        f"(Over: {total_info.get('total_over_money', 'N/A')} | "
                        f"Under: {total_info.get('total_under_money', 'N/A')})"
                    )
            
            game_info.append("=" * 80)
            matches.append("\n".join(game_info))
        
        if matches:
            # Create a detailed response with all match information
            response = [
                f"\n🏀 NBA Schedule for {target_date} - {event_type} 🏀",
                "=" * 80,
                "",  # Empty line for better readability
                "\n\n".join(matches)  # All match details
            ]
            matches_str = "\n".join(response)
            # Print for execution logs and return the same string
            print(f"NBA Matches for {target_date}:\n{matches_str}")
            return matches_str
        else:
            no_matches_msg = f"No NBA matches found with complete information for {target_date}."
            print(no_matches_msg)
            return no_matches_msg
            
    except requests.exceptions.RequestException as e:
        error_msg = f"Error retrieving NBA matches: {str(e)}"
        print(error_msg)
        return error_msg
    except (KeyError, json.JSONDecodeError) as e:
        error_msg = f"Error parsing NBA data: {str(e)}"
        print(error_msg)
        return error_msg

@tool
def predict_nba_match(match_info: str) -> str:
    """
    A tool that generates a prediction for an NBA match.
    
    Args:
        match_info: A string containing match details in the format "TeamA vs TeamB".
    
    Returns:
        A string with the prediction (e.g., "The prediction is that TeamA will win.").
    """
    import random
    teams = match_info.split(" vs ")
    if len(teams) == 2:
        prediction = random.choice(teams)
        return f"The prediction is that {prediction} will win."
    else:
        return "Invalid match format. Please provide details in the format 'TeamA vs TeamB'."

# Instantiate the final_answer tool.
final_answer = FinalAnswerTool()

# Configure the model.
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

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

# Initialize the agent with all necessary tools.
agent = CodeAgent(
    model=model,
    tools=[final_answer, get_current_time_in_timezone, get_nba_matches, predict_nba_match],
    max_steps=6,
    verbosity_level=2,  # Increased verbosity to show more details
    grammar=None,
    planning_interval=None,
    name="NBA Match Assistant",
    description="An assistant that provides detailed NBA match information and predictions",
    prompt_templates=prompt_templates
)

# Use the existing Gradio UI configuration
gradio_ui = GradioUI(agent)
gradio_ui.launch()