import uvicorn
import requests
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Dict
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# API Configuration
TOKEN_URL = "https://spuckhogycrxcbomznwo.supabase.co/auth/v1/token?grant_type=refresh_token"
GPT_URL = "https://not-diamond-workers.t7-cc4.workers.dev/stream-message"

# Global token storage
ACCESS_TOKEN = "eyJhbGciOiJIUzI1NiIsImtpZCI6Im1ZMmorY1dRdWFNTEp3d1AiLCJ0eXAiOiJKV1QifQ.eyJpc3MiOiJodHRwczovL3NwdWNraG9neWNyeGNib216bndvLnN1cGFiYXNlLmNvL2F1dGgvdjEiLCJzdWIiOiI1NzRjZDQ1OS1jNDFiLTQ3MTgtODRlNi0xZWRhNGZiZDI2MWUiLCJhdWQiOiJhdXRoZW50aWNhdGVkIiwiZXhwIjoxNzMyMzY4NDQxLCJpYXQiOjE3MzIzNjQ4NDEsImVtYWlsIjoiYWxob29yc2hvcHBAZ21haWwuY29tIiwicGhvbmUiOiIiLCJhcHBfbWV0YWRhdGEiOnsicHJvdmlkZXIiOiJnb29nbGUiLCJwcm92aWRlcnMiOlsiZ29vZ2xlIl19LCJ1c2VyX21ldGFkYXRhIjp7ImF2YXRhcl91cmwiOiJodHRwczovL2xoMy5nb29nbGV1c2VyY29udGVudC5jb20vYS9BQ2c4b2NMeVN2MUtwakxHMnNTMm44bGtwNzBqSWRid1ZDRTFqWjNEOWQ4c3RSNm5Xb1JVb3c9czk2LWMiLCJlbWFpbCI6ImFsaG9vcnNob3BwQGdtYWlsLmNvbSIsImVtYWlsX3ZlcmlmaWVkIjp0cnVlLCJmdWxsX25hbWUiOiJBbGhvb3IgU2hvcHAiLCJpc3MiOiJodHRwczovL2FjY291bnRzLmdvb2dsZS5jb20iLCJuYW1lIjoiQWxob29yIFNob3BwIiwicGhvbmVfdmVyaWZpZWQiOmZhbHNlLCJwaWN0dXJlIjoiaHR0cHM6Ly9saDMuZ29vZ2xldXNlcmNvbnRlbnQuY29tL2EvQUNnOG9jTHlTdjFLcGpMRzJzUzJuOGxrcDcwaklkYndWQ0UxalozRDlkOHN0UjZuV29SVW93PXM5Ni1jIiwicHJvdmlkZXJfaWQiOiIxMDE3NjY0MjA3Nzk0MjU0OTgwNDYiLCJzdWIiOiIxMDE3NjY0MjA3Nzk0MjU0OTgwNDYifSwicm9sZSI6ImF1dGhlbnRpY2F0ZWQiLCJhYWwiOiJhYWwxIiwiYW1yIjpbeyJtZXRob2QiOiJvYXV0aCIsInRpbWVzdGFtcCI6MTczMjM2MTAxNn1dLCJzZXNzaW9uX2lkIjoiZWZlNTRmMzgtYjFhYS00NjNmLThhODgtM2JiOWVjMTVmYTlkIiwiaXNfYW5vbnltb3VzIjpmYWxzZX0.fCdK8WF4ipmSucvRTtcs0-wgpzkeY6dIlKD62IjYS7M"
REFRESH_TOKEN = "VW2ol3sRN7S_-LqLtBE_fQ"

app = FastAPI(title="AI Response Server")

# Pydantic model for request validation
class MessageRequest(BaseModel):
    messages: List[Dict[str, str]]
    model: str = "chatgpt-4o-latest"
    temperature: float = 0.9

# Token Refresh Function
def refresh_access_token(refresh_token: str):
    global ACCESS_TOKEN
    
    headers = {
        "Authorization": f"Bearer {refresh_token}",
        "Content-Type": "application/json"
    }
    
    payload = {"refresh_token": refresh_token}
    
    try:
        response = requests.post(TOKEN_URL, headers=headers, json=payload)
        response.raise_for_status()
        tokens = response.json()
        
        ACCESS_TOKEN = tokens.get("access_token")
        new_refresh_token = tokens.get("refresh_token")
        
        logger.info("Access token refreshed successfully")
        return ACCESS_TOKEN, new_refresh_token
    
    except requests.exceptions.RequestException as e:
        logger.error(f"Token refresh failed: {e}")
        return None, None

# GPT Message Sending Function
def send_message_to_gpt(messages: List[Dict[str, str]], model: str = "chatgpt-4o-latest", temperature: float = 0.9):
    global ACCESS_TOKEN
    
    headers = {
        "Authorization": f"Bearer {ACCESS_TOKEN}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "messages": messages,
        "model": model,
        "temperature": temperature
    }
    
    try:
        response = requests.post(GPT_URL, headers=headers, json=payload)
        response.raise_for_status()
        
        return response.text
    
    except requests.exceptions.HTTPError as http_err:
        if response.status_code == 401:  # Unauthorized, token might be expired
            logger.warning("Access token expired. Attempting to refresh...")
            new_access_token, _ = refresh_access_token(REFRESH_TOKEN)
            
            if new_access_token:
                # Retry the request with the new token
                return send_message_to_gpt(messages, model, temperature)
        
        logger.error(f"HTTP error occurred: {http_err}")
        raise HTTPException(status_code=500, detail=f"GPT API request failed: {http_err}")
    
    except requests.exceptions.RequestException as req_err:
        logger.error(f"Request error: {req_err}")
        raise HTTPException(status_code=500, detail=f"Request to GPT failed: {req_err}")

# Endpoint for GPT interaction
@app.post("/get-response")
async def get_gpt_response(request: MessageRequest):
    try:
        response = send_message_to_gpt(
            messages=request.messages,
            model=request.model,
            temperature=request.temperature
        )
        return {"response": response}
    except HTTPException as e:
        raise e

# Health check endpoint
@app.get("/health")
async def health_check():
    return {"status": "healthy"}

# DNS Resolution Test Endpoint
@app.get("/test-dns")
async def test_dns():
    try:
        import socket
        socket.gethostbyname('not-diamond-workers.t7-cc4.workers.dev')
        return {"status": "success", "message": "DNS resolution successful"}
    except Exception as e:
        return {"status": "failure", "error": str(e)}

if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host='not-diamond-workers.t7-cc4.workers.dev', port=8000)