from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
from transformers import pipeline
import torch

# Initialize the FastAPI app
app = FastAPI()

# Determine device (use GPU if available, otherwise CPU)
device = 0 if torch.cuda.is_available() else -1

# Initialize the NER pipeline
ner_pipeline = pipeline(
    "ner",
    model="dbmdz/bert-large-cased-finetuned-conll03-english",
    aggregation_strategy="simple",
    device=device
)

# Initialize the QA pipeline
qa_pipeline = pipeline(
    "question-answering",
    model="deepset/roberta-base-squad2",
    device=device
)

# Allowed domains for filtering
allowed_domains = [
    "clothing", "fashion", "shopping", "accessories", "sustainability", "shoes", "hats", "shirts",
    "dresses", "pants", "jeans", "skirts", "jackets", "coats", "t-shirts", "sweaters", "hoodies",
    "activewear", "formal wear", "casual wear", "sportswear", "outerwear", "swimwear", "underwear",
    "lingerie", "socks", "scarves", "gloves", "belts", "ties", "caps", "beanies", "boots", "sandals",
    "heels", "sneakers", "materials", "cotton", "polyester", "wool", "silk", "leather", "denim",
    "linen", "athleisure", "ethnic wear", "fashion trends", "custom clothing", "tailoring",
    "sustainable materials", "recycled clothing", "fashion brands", "streetwear",
    "footwear", "handbags", "jewelry", "watches", "eyewear", "cosmetics", "beauty products",
    "personal care", "fragrances", "home decor", "lifestyle", "luxury goods", "vintage clothing",
    "second-hand clothing", "upcycled fashion", "ethical fashion", "eco-friendly products",
    "fashion technology", "textile innovation", "fashion marketing", "fashion retail"
]

# Context for the QA pipeline
context_msg = (
    "Hutter Products GmbH provides a wide array of services to help businesses create high-quality, sustainable products. "
    "Their offerings include comprehensive product design, ensuring items are both visually appealing and functional, and product consulting, "
    "which provides expert advice on features, materials, and design elements. They also offer sustainability consulting to integrate eco-friendly practices, "
    "such as using recycled materials and Ocean Bound Plastic. Additionally, they manage customized production to ensure products meet the highest standards "
    "and offer product animation services, creating realistic rendered images and animations to enhance online engagement. These services collectively enable "
    "businesses to develop products that are sustainable, market-responsive, and aligned with their brand identity."
)

# Pydantic models for structured responses
class Entity(BaseModel):
    word: str
    entity_group: str
    score: float

class NERResponse(BaseModel):
    entities: List[Entity]
    words: List[str]  # List of extracted words (added)

class QAResponse(BaseModel):
    question: str
    answer: str
    score: float

class CombinedRequest(BaseModel):
    text: str  # The input text prompt

class CombinedResponse(BaseModel):
    ner: NERResponse  # NER output
    qa: QAResponse  # QA output

# Function to check if the input text belongs to allowed domains
def is_text_in_allowed_domain(text: str, domains: List[str]) -> bool:
    for domain in domains:
        if domain in text.lower():
            return True
    return False

# Combined endpoint for NER and QA with domain filtering
@app.post("/process/", response_model=CombinedResponse)
async def process_request(request: CombinedRequest):
    input_text = request.text

    # Check if the input text belongs to the allowed domains
    if not is_text_in_allowed_domain(input_text, allowed_domains):
        raise HTTPException(
            status_code=400,
            detail="The input text does not match the allowed domains. Please provide a query related to clothing, fashion, or accessories."
        )

    # Perform Named Entity Recognition (NER)
    ner_entities = ner_pipeline(input_text)

    # Process NER results into a structured response
    formatted_entities = [
        {
            "word": entity["word"],
            "entity_group": entity["entity_group"],
            "score": float(entity["score"]),
        }
        for entity in ner_entities
    ]
    ner_words = [entity["word"] for entity in ner_entities]  # Collect only the words

    ner_response = {
        "entities": formatted_entities,
        "words": ner_words  # Include the list of words
    }

    # Perform Question Answering (QA)
    qa_result = qa_pipeline(question=input_text, context=context_msg)
    qa_result["score"] = float(qa_result["score"])  # Convert numpy.float32 to Python float

    qa_response = {
        "question": input_text,
        "answer": qa_result["answer"],
        "score": qa_result["score"]
    }

    # Return both NER and QA responses
    return {"ner": ner_response, "qa": qa_response}

# Root endpoint
@app.get("/")
async def root():
    return {"message": "Welcome to the NER and QA API!"}