from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool

import sys  # Add this import to fix the error
import datetime
import time
import requests

from checks.failed_check import  create_failed_gradio_ui
from checks.endpoint_check import check_public_endpoint
from checks.health_check import check_model_endpoint, should_launch_ui

import os
import pytz # Had to give it permission in Code agent  

from tools.tools import load_tools

from agents.model import load_huggingface_model
from agents.prompts import load_prompts

from ContentGradio import ContentAgentUI
   
def initialize_agent(endpoint_uri: str):

     # Load prompts
    combined_prompts = load_prompts()
    combined_tools = load_tools() ## from the tools folder (don't remove final answer)

    
    """Initialize and configure the CodeAgent"""
    model = HfApiModel(
        max_tokens=2096,
        temperature=0.5,
        model_id=endpoint_uri,
        custom_role_conversions=None,
    )
    
    return CodeAgent(
        model=model,
        tools=combined_tools,
        max_steps=6,
        verbosity_level=3,
        name="content_agent",
        description="Evaluates whether text is polite or impolite.",
        prompt_templates=combined_prompts(),
        additional_authorized_imports=["pytz"]
    ) 

def wake_up_endpoint(endpoint_uri, max_wait=30):
    """Poll the endpoint until it responds OK or timeout."""

    # Get token from environment variable or secrets store
    hf_token = os.environ.get("HF_TOKEN")  # Or replace with your secret loading logic
    headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
    
    payload = {"inputs": "ping"}  # Adjust to minimal valid input
    start = time.time()

    while time.time() - start < max_wait:
        try:
            print("Pinging endpoint...")
            print(endpoint_uri)
            response = requests.post(endpoint_uri, json=payload, timeout=5)
            
            if response.ok:
                print("✅ Endpoint is awake.")
                return True

            # If it's a 503 or 504, it's likely still warming up
            print(f"Response: {response.status_code}, retrying...")
        except requests.exceptions.RequestException as e:
            print(f"Exception: {e}, retrying...")

        time.sleep(3)  # fixed short backoff; can be exponential if needed

    print("❌ Timed out waiting for endpoint to wake up.")
    return False

def main():
    # Load endpoint and check health

    # Prechecks 
    endpoint_uri = load_huggingface_model() # Get the URI for the endpoint 
   

    # Wake it up before health check
    wake_up_successful = wake_up_endpoint(endpoint_uri)
    if not wake_up_successful:
        print("Warning: Could not wake up the endpoint. Exiting.")
        exit(0)

    is_healthy, status_info = check_model_endpoint(endpoint_uri) # Test the endpoint
    
    if not is_healthy:
        interface = create_failed_gradio_ui(status_info)
        interface.launch(show_error=True, share=True)
        return
    
    # Initialize and run the agent
    agent = initialize_agent(endpoint_uri)
    # Create an instance of the ContentAgentUI class
    ui = ContentAgentUI()
    ui.pass_through_agent(agent)  # Pass through the agent
    

if __name__ == "__main__":
    main()