import os
from openai import OpenAI
from models.input import Input
import re
from transformers import pipeline
from services.utils import clean_text

intent_recognizer = pipeline(
    "zero-shot-classification", model="facebook/bart-large-mnli")

ner_recognizer = pipeline(
    'ner', model='dbmdz/bert-large-cased-finetuned-conll03-english')

openai_key = os.environ.get("OPENAI_KEY")
openai_client = OpenAI(api_key=openai_key)

transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-base.en")

# Define regex patterns for entities
patterns = {
    'Product': r'\b(iphone|samsung|macbook|ps5|galaxy|pixel|shoes|shampoo|cellphone|smartphone|tablet|laptop|headphones|console|tv|camera)\b',
    'Brand': r'\b(apple|samsung|google|sony|microsoft|dell|hp|lenovo|asus|nintendo|canon|nikon)\b',
    'Category': r'\b(laptops|dresses|phones|electronics|clothing|footwear|accessories|home appliances|furniture)\b',
    'Color': r'\b(red|black|yellow|blue|green|white|grey|pink|purple|orange|brown)\b',
    'Price Range': r'\b(under \$?\d+|below \$?\d+|less than \$?\d+|more than \$?\d+|above \$?\d+|between \$?\d+ and \$?\d+)\b',
    'Quantity': r'\b(\d+ bottles|\d+ items|\d+ pieces|\d+ units|\d+)\b',
    'Order Number': r'\bB-\d+\b',
    'Issue': r'\b(account help|payment issue|order problem|shipping delay|return request|product complaint|technical support)\b'
}

INTENTS = [
    "search for products",
    "order management",
    "checkout",
    "customer support",
    "account management"
]

def transcribe(audio):
    sr, y = audio
    y = y.astype(np.float32)
    y /= np.max(np.abs(y))
    return transcriber({"sampling_rate": sr, "raw": y})["text"]  


def recognize_intent(text):
    cleaned_text = clean_text(text)
    intent = intent_recognizer(cleaned_text, INTENTS)
    return intent['labels'][0]


def recognize_entities(text):
    cleaned_text = clean_text(text)
    entities_from_ner = ner_recognizer(cleaned_text)
    entities_from_re = {entity: re.findall(pattern, text.lower(
    )) for entity, pattern in patterns.items() if re.findall(pattern, text.lower())}

    return entities_from_re


def extract_keywords(text):
    cleaned_text = clean_text(text)
    return cleaned_text.split()


def generate_response(input: Input, context_from_elasticsearch):
    prompt = input.history
    prompt.append({"role": "assistant", "content": str(context_from_elasticsearch)})

    print(prompt)

    response = openai_client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=prompt
    )

    return response.choices[0].message.content