import json
import os

import google.generativeai as genai
from dotenv import load_dotenv
from pydantic import BaseModel

load_dotenv()

GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
genai.configure(api_key=GOOGLE_API_KEY)


class OutputSchema(BaseModel):
    topic: str
    answer: str


def valid_output(output: dict) -> dict:
    return OutputSchema(**output)


class GenAI:
    def __init__(self):
        self.genai_model = genai.GenerativeModel("gemini-2.0-flash-lite")

    def generate_content(
        self,
        question: str,
        question_context: str,
        unique_topics: str,
        few_shot_examples: str,
    ):
        prompt = f"""
    INSTRUCTIONS:
        1. You are a expert assistant to mental health counselors.
        2. You are given following:
            2.1 Input Question: A input question from mental health counselor seeking assistance.
            2.2 Input Question Context: Additional context for input question. Can be None. If available, utilize this context in your response.
            2.3 Topics: Categories of topics to which any input question may belong. Any input question will be categorized to one of these topics.
            2.4 Few-shot examples: Some examples
        3. Output the following:
            3.1 A topic from list of topics for input question.
            3.2 A precise answer for the input question.
            - Length of answer should not exceed 256 words.
        4. Your output MUST be a VALID JSON format.
            4.1 Follow the sample JSON format given below.

    INPUT:
        Input Question: {question}
        Input Question Context: {question_context}

        Topics: {unique_topics}

        Few-shot examples: {few_shot_examples}

    Sample output JSON format: [{{
        "topic": "A topic from list of topics for input question."
        "answer": "An answer for the input question."
        }}]

    OUTPUT:"""
        response = self.genai_model.generate_content(prompt)
        out = response.text
        out = out[7:-3]
        out = json.loads(out)
        valid_response = valid_output(out[0])

        return (valid_response.answer, valid_response.topic)