# === Imports ===
import os
import time
import requests
import socket
from dotenv import load_dotenv
from supabase import create_client
from sentence_transformers import SentenceTransformer
from openai import OpenAI

# === Load Environment Variables ===
load_dotenv()

# === OpenAI Configuration ===
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise ValueError("OPENAI_API_KEY is not set in the environment variables.")
openai_client = OpenAI(api_key=OPENAI_API_KEY)

# === Supabase Configuration ===
SUPABASE_URL = "https://lmpazoxzucnlqqxjoihi.supabase.co"
SUPABASE_KEY = os.getenv("SUPABASE_API_KEY")
if not SUPABASE_KEY:
    raise ValueError("SUPABASE_KEY is not set in the environment variables.")
supabase = create_client(SUPABASE_URL, SUPABASE_KEY)

# === Embedding Model for Scoring ===
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# === Hugging Face API Configuration ===
HF_API_TOKEN = os.getenv("HF_API_TOKEN")
if not HF_API_TOKEN:
    raise ValueError("Missing Hugging Face API key. Check your .env file.")

# Headers for API requests
HF_HEADERS = {"Authorization": f"Bearer {HF_API_TOKEN}"}

# === Hugging Face Model Endpoints ===
HF_ENDPOINTS = {
    "bart-large-cnn-ovt": {
        "url": "https://hedemwou4oqkk65c.us-east-1.aws.endpoints.huggingface.cloud",
        "task": "summarization",
        "model_id": "facebook/bart-large-cnn"
    },
    "vzwjawyxvu030jsw": {  # Updated endpoint name to match URL
        "url": "https://vzwjawyxvu030jsw.us-east-1.aws.endpoints.huggingface.cloud",
        "task": "text-generation",
        "model_id": "google/gemma-7b"
    }
}

def check_endpoint_status(endpoint_name: str) -> dict:
    """
    Check the status of a private Hugging Face endpoint using DNS resolution
    """
    if endpoint_name not in HF_ENDPOINTS:
        return {
            "status": "error",
            "error": f"Unknown endpoint: {endpoint_name}"
        }

    try:
        endpoint_info = HF_ENDPOINTS[endpoint_name]
        hostname = endpoint_info['url'].replace('https://', '').split('/')[0]
        
        # Try DNS resolution
        try:
            socket.gethostbyname(hostname)
            # If DNS resolves, endpoint exists but may be stopped
            return {
                "status": "stopped",
                "scaled": True,
                "pending": 0,
                "error": None
            }
        except socket.gaierror:
            # If DNS fails, endpoint doesn't exist
            return {
                "status": "error",
                "error": "Endpoint not found"
            }
    except Exception as e:
        return {
            "status": "error",
            "error": str(e)
        }

def toggle_endpoint(endpoint_name: str, action: str) -> dict:
    """
    Start or stop a private Hugging Face endpoint
    """
    try:
        # For private endpoints, use the Endpoints API
        api_base = "https://api.endpoints.huggingface.cloud"
        action_url = f"{api_base}/v2/endpoint/{endpoint_name}/{action}"
        
        response = requests.post(
            action_url,
            headers=HF_HEADERS,
            timeout=10
        )
        
        if response.status_code in [200, 202]:
            return {
                "success": True,
                "message": f"Successfully {action}ed endpoint"
            }
        else:
            return {
                "error": f"Failed to {action} endpoint: {response.text}"
            }
    except Exception as e:
        return {
            "error": f"Failed to {action} endpoint: {str(e)}"
        }

# === Query Helper ===
def query(payload: dict, endpoint_name: str) -> dict:
    """
    Send a query to a Hugging Face endpoint
    """
    if endpoint_name not in HF_ENDPOINTS:
        return {
            "error": f"Unknown endpoint: {endpoint_name}"
        }

    endpoint_info = HF_ENDPOINTS[endpoint_name]
    url = endpoint_info['url']

    try:
        response = requests.post(
            url,
            headers=HF_HEADERS,
            json=payload,
            timeout=30
        )

        if response.status_code == 200:
            return response.json()
        else:
            return {
                "error": f"Query failed with status {response.status_code}: {response.text}"
            }
    except Exception as e:
        return {
            "error": str(e)
        }