import os
import gc
import glob
import time
import logging
from typing import List, Union, Optional, Literal
from fastapi import FastAPI, Depends, HTTPException, Header, status
from fastapi.responses import JSONResponse, FileResponse
from pydantic import BaseModel, Field
from llama_cpp import Llama

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

# 🔒 SEGURANÇA: Falha imediata se API_KEY não estiver configurada
API_KEY = os.getenv("API_KEY")
if not API_KEY:
    raise RuntimeError(
        "❌ SECURITY ERROR: API_KEY environment variable is not set!\n"
        "Please configure the API_KEY secret in your Hugging Face Space:\n"
        "Settings → Variables and Secrets → Add secret\n"
        "Name: API_KEY | Value: [your-secure-key]"
    )

MODEL_PATH = glob.glob("/app/model/*.gguf")[0]
logger.info(f"Carregando modelo: {MODEL_PATH}")

load_start = time.time()
llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=512,
    n_threads=2,
    n_batch=512,
    embedding=True,
    verbose=False
)
load_time = time.time() - load_start
logger.info(f"🚀 Model loaded in {load_time:.2f}s")

app = FastAPI(title="Embedding Engine (GGUF Q8_0)")

@app.get("/", include_in_schema=False)
async def root():
    return FileResponse("index.html")

class EmbeddingRequest(BaseModel):
    input: Union[str, List[str]]
    model: str = "LiquidAI/LFM2.5-Embedding-350M"
    input_type: Optional[Literal["query", "document"]] = Field(
        default="document",
        alias="inputType"
    )
    model_config = {"populate_by_name": True}

class EmbeddingObject(BaseModel):
    object: str = "embedding"
    embedding: List[float]
    index: int

class EmbeddingResponse(BaseModel):
    object: str = "list"
    data: List[EmbeddingObject]
    model: str
    usage: dict

def verify_api_key(authorization: str = Header(None)):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or missing Authorization header",
            headers={"WWW-Authenticate": "Bearer"},
        )
    token = authorization.replace("Bearer ", "")
    if token != API_KEY:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid API key",
            headers={"WWW-Authenticate": "Bearer"},
        )

@app.post("/v1/embeddings", response_model=EmbeddingResponse)
async def create_embeddings(
    request: EmbeddingRequest,
    _: str = Depends(verify_api_key)
):
    req_start = time.time()
    
    inputs = request.input if isinstance(request.input, list) else [request.input]
    
    if request.input_type == "query":
        inputs = [f"query: {text}" for text in inputs]
    else:
        inputs = [f"document: {text}" for text in inputs]
    
    prep_time = time.time() - req_start
    logger.info(f"📝 Prep time: {prep_time:.3f}s | Input count: {len(inputs)} | First input length: {len(inputs[0])} chars")
    
    try:
        infer_start = time.time()
        response = llm.create_embedding(input=inputs)
        infer_time = time.time() - infer_start
        
        logger.info(f"🧠 Inference time: {infer_time:.3f}s")
        
        response["model"] = request.model
        
        total_time = time.time() - req_start
        logger.info(f"✅ Total request time: {total_time:.3f}s")
        
        return JSONResponse(content=response)
    except Exception as e:
        logger.error(f"Embedding error: {str(e)}")
        raise HTTPException(status_code=500, detail="Embedding generation failed")
    finally:
        gc.collect()