from langchain import hub, callbacks
from langchain_core.output_parsers import JsonOutputParser
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, UUID4
from .suggest_expectations import Expectation
from typing import List
import os

class Inputs(BaseModel):
    course: str
    module: str
    concepts: List[str]
    expectation: str
    check_question: str
    request: str

class Response(BaseModel):
    run_id: UUID4
    expectation: str
    check_question: str

class ExpectationRevision:
    def kickoff(self, inputs={}):
        self.learning_outcome = inputs["expectation"]
        self.check_question = inputs["check_question"]
        self.request = inputs["request"]
        self.course = inputs["course"]
        self.module = inputs["module"]
        self.concepts =  inputs["concepts"]
        llm_response = self._get_suggestion()
        return {
            "run_id": self.run_id,
            "expectation": llm_response["expectation"],
            "check_question": llm_response["check_question"]
        }

    def _get_suggestion(self):
        parser = JsonOutputParser(pydantic_object=Expectation)
        prompt = hub.pull("course_learn_expectation_revision")
        llm = ChatOpenAI(model=os.environ['OPENAI_MODEL'], temperature=0.2)
        chain = (prompt | llm | parser).with_config({
            "tags": ["course_learn", "suggest_expectations"], "run_name": "Suggest Module Expectations",
            "metadata" : {
                "version": "v1.0.0",
                "growth_activity": "course_learn",
                "env": os.environ["ENV"],
                "model": os.environ["OPENAI_MODEL"],
            }
        })

        with callbacks.collect_runs() as cb:
            response = chain.invoke({
                "learning_outcome": self.learning_outcome, "check_question": self.check_question, "request": self.request,
                "course": self.course, "module": self.module, "concepts": "* " + ("\n* ".join(self.concepts)),
                "format_instructions": parser.get_format_instructions()
              })
            self.run_id = cb.traced_runs[0].id

        return response


# Example usage
# rework = ExpectationRevision()
# response = rework.kickoff(inputs={
#     "expectation": "Recognize the importance of query rewriting and how to transform inefficient queries into more efficient ones.",
#     "check_question": "Can you provide an example of a poorly written SQL query and demonstrate how you would rewrite it to optimize its performance?",
#     "request": "Can you provide the poorly written SQL query that user has to rewrite"
# })
# print(response)