import os
import requests
from fastapi import FastAPI, Request, Depends
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import Optional, Any

# Check whether we are executing inside a Hugging Face Space
SPACE_NAME = os.getenv("SPACE_NAME", default=None)
if SPACE_NAME is not None:
    print(f"Running inside {SPACE_NAME} Space.")
    try:
        # Try to auto-login using the Space's environment variables
        login(automatically=True)
    except Exception as e:
        print(f"Failed to auto-login ({str(e)}). Manually check the HF_ACCESS_TOKEN environment variable.")
        sys.exit(1)

try:
    HUGGINGFACE_TOKEN = os.environ['HF_ACCESS_TOKEN']
except KeyError:
    print('The environment variable "HF_ACCESS_TOKEN" is not found. Please configure it correctly in your Space.')
    sys.exit(1)

# Set up the API endpoint and headers
model_id = "152334H/miqu-1-70b-sf"
endpoint = f"https://api-inference.huggingface.co/models/{model_id}"
headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"}

# App definition
app = FastAPI()

# Helper function to read raw request bodies
async def parse_raw(request: Request):
    return await request.body()

# Generate text using the Inference API
def generate_text(prompt: str) -> str:
    data = {
        "inputs": prompt,
        "options": {
            "max_new_tokens": 200,
            "temperature": 0.7,
            "top_p": 0.95,
            "use_cache": False,
        },
    }

    response = requests.post(endpoint, headers=headers, json=data)
    return response.json()["generated_text"]

# Route for generating text
@app.post("/generate_text")
async def generate_text_route(data: BaseModel = Depends(parse_raw)):
    input_text = data.raw.decode("utf-8")
    if not input_text or len(input_text) <= 0:
        return JSONResponse({"error": "Empty input received."}, status_code=400)

    return {"output": generate_text(input_text)}

# Mount static files
app.mount("/static", StaticFiles(directory="static"), name="static")