"""
GLiNER2 FastAPI Server - REST API for entity extraction
"""

import os
import time
from typing import Dict, List, Optional

import torch
import uvicorn
from fastapi import Depends, FastAPI, HTTPException, Security
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from gliner2 import GLiNER2
from pydantic import BaseModel

# Configuration
MODEL_NAME = os.environ.get("GLINER_MODEL", "fastino/gliner2-base-v1")
API_KEY = os.environ.get("GLINER_API_KEY")

app = FastAPI(
    title="GLiNER2 API",
    description="GPU-accelerated Named Entity Recognition",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

security = HTTPBearer()

if not API_KEY:
    print("WARNING: GLINER_API_KEY not set. API will accept any token.")


def verify_token(credentials: HTTPAuthorizationCredentials = Security(security)):
    """Verify the Bearer token"""
    if API_KEY and credentials.credentials != API_KEY:
        raise HTTPException(status_code=401, detail="Unauthorized")
    return True


# Global model state
model = None
device = None


class ExtractRequest(BaseModel):
    text: str
    labels: List[str]
    threshold: Optional[float] = 0.3


class ExtractResponse(BaseModel):
    entities: Dict[str, List[str]]
    inference_time: float
    device: str


class HealthResponse(BaseModel):
    status: str
    model: str
    device: str
    gpu_available: bool
    gpu_name: Optional[str] = None


@app.on_event("startup")
async def load_model():
    """Load GLiNER2 model on startup"""
    global model, device

    print(f"Loading GLiNER2 model: {MODEL_NAME}")
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    model = GLiNER2.from_pretrained(MODEL_NAME)
    model = model.to(device)
    model.eval()

    print(f"Model loaded on {device}")
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9
        print(f"GPU: {gpu_name} ({gpu_memory:.1f} GB)")


@app.get("/health", response_model=HealthResponse)
async def health():
    """Health check endpoint"""
    gpu_name = None
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)

    return HealthResponse(
        status="running",
        model=MODEL_NAME,
        device=device,
        gpu_available=torch.cuda.is_available(),
        gpu_name=gpu_name,
    )


@app.post("/extract", response_model=ExtractResponse)
async def extract_entities(
    request: ExtractRequest,
    authorized: bool = Depends(verify_token),
):
    """Extract entities from text"""
    start_time = time.time()
    result = model.extract_entities(
        request.text,
        request.labels,
        threshold=request.threshold,
    )
    inference_time = time.time() - start_time

    return ExtractResponse(
        entities=result.get("entities", {}),
        inference_time=inference_time,
        device=device,
    )


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")

import requests

API_URL = "http://<IP>:<PORT>"

def extract_entities(text, labels, threshold=0.3):
    response = requests.post(
        f"{API_URL}/extract",
        headers={"Authorization": "Bearer gliner-api-key"},
        json={"text": text, "labels": labels, "threshold": threshold}
    )
    response.raise_for_status()
    return response.json()

result = extract_entities(
    text="Elon Musk announced Tesla's new factory in Berlin.",
    labels=["person", "company", "location"]
)
print(result["entities"])
# {'person': ['Elon Musk'], 'company': ['Tesla'], 'location': ['Berlin']}