import datetime
import requests
import pytz
import yaml
import os
from typing import Dict, Union

from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Load the API key from Hugging Face Spaces secrets (environment variable)
API_KEY = os.getenv("CRICKETDATA_API_KEY", "default_key_if_not_found")
BASE_URL = "https://api.cricketdata.org/v1"

# Custom Tool: Analyze Cricketer Form
@tool
def analyze_cricketer_form(player_name: str, role: str) -> str:
    """A tool to analyze a cricketer's recent form based on their role and name.
    
    Args:
        player_name: The name of the cricketer (e.g., 'Virat Kohli').
        role: The player's role ('batter', 'bowler', 'wicket_keeper', 'fielder').
    
    Returns:
        A string summary of the player's career stats and recent form.
    """
    valid_roles = ['batter', 'bowler', 'wicket_keeper', 'fielder']
    if role.lower() not in valid_roles:
        return f"Error: Invalid role '{role}'. Choose from {valid_roles}."

    player_id = get_player_id(player_name)
    if not player_id:
        return f"Error: Player '{player_name}' not found."

    career_stats = get_career_stats(player_id)
    if not career_stats:
        return f"Error: Could not fetch career stats for '{player_name}'."

    recent_matches = get_recent_matches(player_id)
    
    response = f"Analysis of {player_name} ({role.capitalize()}):\n\n"
    response += "Career Stats:\n" + format_career_stats(career_stats, role) + "\n\n"
    response += "Recent Form (Last 5 Matches):\n" + analyze_recent_form(recent_matches, role, player_name)
    
    return response

def get_player_id(player_name: str) -> Union[str, None]:
    url = f"{BASE_URL}/players/search?query={player_name}&apikey={API_KEY}"
    try:
        response = requests.get(url)
        response.raise_for_status()
        data = response.json()
        if data.get("data") and len(data["data"]) > 0:
            return data["data"][0]["id"]
        return None
    except requests.RequestException as e:
        print(f"Error fetching player ID: {e}")
        return None

def get_career_stats(player_id: str) -> Union[Dict, None]:
    url = f"{BASE_URL}/players/{player_id}/stats?apikey={API_KEY}"
    try:
        response = requests.get(url)
        response.raise_for_status()
        return response.json()
    except requests.RequestException as e:
        print(f"Error fetching career stats: {e}")
        return None

def get_recent_matches(player_id: str) -> list:
    url = f"{BASE_URL}/players/{player_id}/matches?apikey={API_KEY}"
    try:
        response = requests.get(url)
        response.raise_for_status()
        matches = response.json().get("data", [])
        return matches[:5]  # Limit to last 5 matches
    except requests.RequestException as e:
        print(f"Error fetching recent matches: {e}")
        return []

def format_career_stats(stats: Dict, role: str) -> str:
    if not stats or "data" not in stats:
        return "No stats available."
    
    data = stats["data"]
    output = ""
    
    if role == "batter":
        batting = data.get("batting", {})
        output += f"Matches: {batting.get('matches', 'N/A')}\n"
        output += f"Runs: {batting.get('runs', 'N/A')}\n"
        output += f"Average: {batting.get('average', 'N/A')}\n"
        output += f"Strike Rate: {batting.get('strike_rate', 'N/A')}\n"
    
    elif role == "bowler":
        bowling = data.get("bowling", {})
        output += f"Matches: {bowling.get('matches', 'N/A')}\n"
        output += f"Wickets: {bowling.get('wickets', 'N/A')}\n"
        output += f"Economy: {bowling.get('economy', 'N/A')}\n"
        output += f"Best: {bowling.get('best', {}).get('figures', 'N/A')}\n"
    
    elif role == "wicket_keeper":
        fielding = data.get("fielding", {})
        output += f"Catches: {fielding.get('catches', 'N/A')}\n"
        output += f"Stumpings: {fielding.get('stumpings', 'N/A')}\n"
    
    elif role == "fielder":
        fielding = data.get("fielding", {})
        output += f"Catches: {fielding.get('catches', 'N/A')}\n"
        output += f"Run Outs: {fielding.get('run_outs', 'N/A')}\n"
    
    return output

def analyze_recent_form(matches: list, role: str, player_name: str) -> str:
    if not matches:
        return "No recent match data available."
    
    output = ""
    if role == "batter":
        total_runs = 0
        games = len(matches)
        for match in matches:
            batting = match.get("batting", {})
            runs = batting.get("runs", 0)
            total_runs += runs
            output += f"Match vs {match.get('opponent', 'Unknown')}: {runs} runs\n"
        avg = total_runs / games if games > 0 else 0
        output += f"Recent Average: {avg:.2f}\n"
        output += f"Form Comment: {'Strong' if avg > 30 else 'Inconsistent'} batting form."
    
    elif role == "bowler":
        total_wickets = 0
        games = len(matches)
        for match in matches:
            bowling = match.get("bowling", {})
            wickets = bowling.get("wickets", 0)
            total_wickets += wickets
            output += f"Match vs {match.get('opponent', 'Unknown')}: {wickets} wickets\n"
        avg_wickets = total_wickets / games if games > 0 else 0
        output += f"Recent Wickets per Game: {avg_wickets:.2f}\n"
        output += f"Form Comment: {'Strong' if avg_wickets > 1 else 'Inconsistent'} bowling form."
    
    elif role == "wicket_keeper":
        total_stumpings = sum(m.get("fielding", {}).get("stumpings", 0) for m in matches)
        total_catches = sum(m.get("fielding", {}).get("catches", 0) for m in matches)
        output += f"Stumpings: {total_stumpings}, Catches: {total_catches}\n"
        output += f"Form Comment: {'Reliable' if total_stumpings + total_catches > 3 else 'Average'} keeping."
    
    elif role == "fielder":
        total_catches = sum(m.get("fielding", {}).get("catches", 0) for m in matches)
        total_run_outs = sum(m.get("fielding", {}).get("run_outs", 0) for m in matches)
        output += f"Catches: {total_catches}, Run Outs: {total_run_outs}\n"
        output += f"Form Comment: {'Sharp' if total_catches + total_run_outs > 3 else 'Average'} fielding."
    
    return output

# Example tool from your template
@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:
        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)}"

# Initialize tools
final_answer = FinalAnswerTool()
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# Model setup
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
agent = CodeAgent(
    model=model,
    tools=[final_answer, analyze_cricketer_form, get_current_time_in_timezone, image_generation_tool],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    prompt_templates=prompt_templates
)

# Launch the Gradio UI
if __name__ == "__main__":
    GradioUI(agent).launch()