import os
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

# Configurar la caché en un directorio accesible
os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"

# Inicializa la aplicación FastAPI
app = FastAPI()

# Modelo Hugging Face para NER
MODEL_NAME = "Davlan/bert-base-multilingual-cased-ner-hrl"  # Modelo para NER multilingüe
ner_pipeline = pipeline("ner", model=MODEL_NAME)

# Define el esquema de entrada
class JobOffer(BaseModel):
    texto: str

@app.get("/")
def read_root():
    """
    Endpoint de prueba para confirmar que el servidor está activo.
    """
    return {"message": "API funcionando correctamente. Usa /procesar_oferta/ para enviar datos."}


@app.post("/procesar_oferta")
async def procesar_oferta(oferta: JobOffer):
    """
    Procesa el texto de una oferta laboral y genera un JSON estructurado.
    """
    # Aplica el modelo para extraer entidades
    entities = ner_pipeline(oferta.texto)

    # Inicializa el resultado en formato JSON
    resultado = {
        "titulo": None,
        "empresa": None,
        "ubicacion": None,
        "salario": None,
        "expe_min": None,
        "tipo_contrato": None,
        "estudios_min": None,
        "skills": [],
        "descripcion": None,
    }

    # Procesa las entidades para rellenar las claves del JSON
    for entity in entities:
        palabra = entity['word']
        etiqueta = entity['entity']

        if etiqueta == "B-TITLE":
            resultado["titulo"] = palabra if not resultado["titulo"] else f"{resultado['titulo']} {palabra}"
        elif etiqueta == "B-ORG":
            resultado["empresa"] = palabra if not resultado["empresa"] else f"{resultado['empresa']} {palabra}"
        elif etiqueta == "B-LOC":
            resultado["ubicacion"] = palabra if not resultado["ubicacion"] else f"{resultado['ubicacion']} {palabra}"
        elif etiqueta == "B-SALARY":
            resultado["salario"] = palabra if not resultado["salario"] else f"{resultado['salario']} {palabra}"
        elif etiqueta == "B-EXP":
            resultado["expe_min"] = palabra if not resultado["expe_min"] else f"{resultado['expe_min']} {palabra}"
        elif etiqueta == "B-CONTRACT":
            resultado["tipo_contrato"] = palabra if not resultado["tipo_contrato"] else f"{resultado['tipo_contrato']} {palabra}"
        elif etiqueta == "B-STUDIES":
            resultado["estudios_min"] = palabra if not resultado["estudios_min"] else f"{resultado['estudios_min']} {palabra}"
        elif etiqueta == "B-SKILL":
            resultado["skills"].append(palabra)
        elif etiqueta == "B-DESCRIPTION":
            resultado["descripcion"] = palabra if not resultado["descripcion"] else f"{resultado['descripcion']} {palabra}"

    return resultado