# Warning control
import warnings
warnings.filterwarnings('ignore')

from typing import List
import json
from pydantic import BaseModel
# from langchain_community.llms import HuggingFaceHub
from langchain_huggingface import HuggingFaceEndpoint
from crewai import Agent, Task, Crew

import os
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv()) # read local .env file
# hf_api_key = os.environ['HF_API_KEY']
hf_api_key = os.getenv('HF_API_KEY')

llm = HuggingFaceEndpoint(
    repo_id="HuggingFaceH4/zephyr-7b-beta",
    huggingfacehub_api_token=hf_api_key,
    task="text-generation"
)

graphicDesigner = Agent(
    role="Graphic Designer",
    goal="Provide the list of relevant questions that can be ask to user based on his/her requirement: {requirement}",
    backstory="You're working as expert graphic designer "
              "A user has shared with you the requirement: {requirement}."
              "Use information present in requirement"
              "and provide the list of questions that can be asked to user"
              "to gather more specific information based on requirement",
    llm=llm,
    allow_delegation=False,
	verbose=False
)

class Questions(BaseModel):
    QuestionsList: List[str]

Ask_questions = Task(
    description=(
        "Provide the list of relevant questions that can be ask to user to gather more specific information based on requirement: {requirement}"
    ),
    expected_output="List of relevant questions based on user requirement. output only questions",
    # output_json=Questions,
    agent=graphicDesigner,
)

crew = Crew(
    agents=[graphicDesigner],
    tasks=[Ask_questions],
    verbose=False
)

def questions(requirement, userDefinedQuestions=None):
    # "I want a logo for my new business."

    if not userDefinedQuestions:
        result = crew.kickoff(inputs={"requirement": requirement})
        result = result.split('\n')
        # json_dict = json.loads(result)
        questions = [i.split('. ')[1] for i in result]
        return questions
    else:
        return userDefinedQuestions