import json

import google.generativeai as genai
from pydantic import BaseModel, Field

from document_ai_agents.schema_utils import prepare_schema_for_gemini


class Answer(BaseModel):
    answer: str = Field(..., description="Your Answer.")


model = genai.GenerativeModel("gemini-1.5-flash-002")

answer_schema = prepare_schema_for_gemini(Answer)


question = "List all the reasons why LLM hallucinate"

context = (
    "LLM hallucination refers to the phenomenon where large language models generate plausible-sounding but"
    " factually incorrect or nonsensical information. This can occur due to various factors, including biases"
    " in the training data, the inherent limitations of the model's understanding of the real world, and the "
    "model's tendency to prioritize fluency and coherence over accuracy."
)

messages = (
    [context]
    + [
        f"Answer this question: {question}",
    ]
    + [
        f"Use this schema for your answer: {answer_schema}",
    ]
)

response = model.generate_content(
    messages,
    generation_config={
        "response_mime_type": "application/json",
        "response_schema": answer_schema,
        "temperature": 0.0,
    },
)

response = Answer(**json.loads(response.text))

print(f"{response.answer=}")

class AnswerChainOfThoughts(BaseModel):
    rationale: str = Field(
        ...,
        description="Justification of your answer.",
    )
    answer: str = Field(
        ..., description="Your Answer. Answer with 'N/A' if answer is not found"
    )

def answer_question(self, state: DocumentQAState):
        logger.info(f"Responding to question '{state.question}'")
        assert (
            state.pages_as_base64_jpeg_images or state.pages_as_text
        ), "Input text or images"
        messages = (
            [
                {"mime_type": "image/jpeg", "data": base64_jpeg}
                for base64_jpeg in state.pages_as_base64_jpeg_images
            ]
            + state.pages_as_text
            + [
                f"Answer this question: {state.question}",
            ]
            + [
                f"Use this schema for your answer: {self.answer_cot_schema}",
            ]
        )

        response = self.model.generate_content(
            messages,
            generation_config={
                "response_mime_type": "application/json",
                "response_schema": self.answer_cot_schema,
                "temperature": 0.0,
            },
        )

        answer_cot = AnswerChainOfThoughts(**json.loads(response.text))

        return {"answer_cot": answer_cot}

def reformulate_answer(self, state: DocumentQAState):
        logger.info("Reformulating answer")
        if state.answer_cot.answer == "N/A":
            return

        messages = [
            {
                "role": "user",
                "parts": [
                    {
                        "text": "Reformulate this question and its answer as a single assertion."
                    },
                    {"text": f"Question: {state.question}"},
                    {"text": f"Answer: {state.answer_cot.answer}"},
                ]
                + [
                    {
                        "text": f"Use this schema for your answer: {self.declarative_answer_schema}"
                    }
                ],
            }
        ]

        response = self.model.generate_content(
            messages,
            generation_config={
                "response_mime_type": "application/json",
                "response_schema": self.declarative_answer_schema,
                "temperature": 0.0,
            },
        )

        answer_reformulation = AnswerReformulation(**json.loads(response.text))

        return {"answer_reformulation": answer_reformulation}

 def verify_answer(self, state: DocumentQAState):
        logger.info(f"Verifying answer '{state.answer_cot.answer}'")
        if state.answer_cot.answer == "N/A":
            return
        messages = [
            {
                "role": "user",
                "parts": [
                    {
                        "text": "Analyse the following context and the assertion and decide whether the context "
                        "entails the assertion or not."
                    },
                    {"text": f"Context: {state.answer_cot.relevant_context}"},
                    {
                        "text": f"Assertion: {state.answer_reformulation.declarative_answer}"
                    },
                    {
                        "text": f"Use this schema for your answer: {self.verification_cot_schema}. Be Factual."
                    },
                ],
            }
        ]
    
        response = self.model.generate_content(
            messages,
            generation_config={
                "response_mime_type": "application/json",
                "response_schema": self.verification_cot_schema,
                "temperature": 0.0,
            },
        )
    
        verification_cot = VerificationChainOfThoughts(**json.loads(response.text))
    
        return {"verification_cot": verification_cot}