from typing import List, Iterable, Optional
from pydantic import BaseModel, ValidationInfo, model_validator, Field,  field_validator
import instructor
import openai
import asyncio
import os


from groq import AsyncGroq
# Initialize with API key
client = AsyncGroq(api_key=os.getenv("GROQ_API_KEY"))

# Enable instructor patches for Groq client
client = instructor.from_groq(client)

"""
client = instructor.from_openai(
    openai.AsyncOpenAI(
        base_url="http://localhost:11434/v1",
        api_key="ollama",
    ),
    mode=instructor.Mode.JSON,
)
"""
llm = 'llama-3.1-8b-instant' if os.getenv("GROQ_API_KEY") else "qwen2.5" #"gemma3:12b" #"llama3.2" #"deepseek-r1"

class Judgment(BaseModel):
    thought: str = Field(...,
        description="The step-by-step reasoning process used to analyze the reasoning and the answer", examples=["Let's think step by step...context explicit stated donation, therefore answer should be donation"]
    )
    justification: str = Field(...,
        description="Explanation for the logical judgment, detailing key factors that led to the conclusion", examples=["sound reasoning if context stated donation, it is valid and logical for answer to be donation"]
    )
    logical: bool = Field(...,
        description="Boolean judgment indicating whether the reasoning and the answer are logical or relevant (True) or not (False)", examples=[True, False]
    )

#MaybeJudgment=instructor.Maybe(Judgment)
async def judge_relevance(reasoning: str, context:str, answer: str) -> Judgment:
    return await client.chat.create(
        model=llm, #"gpt-4",
        temperature=0.3,
        max_retries=3,
        messages=[
            {
                "role": "system",
                "content": """
                    You are tasked with comparing a (reasoning + context) and a answer to determine if they are relevant to each other or logical in some way. Your goal is to analyze the content, context, and potential connections between the two.

                    To determine if the (reasoning + context) and answer are relevant or logical, please follow these steps:

                    1. Carefully read and understand both the reasoning and the answer.
                    2. Identify the main topic, keywords, and concepts in the (reasoning + context).
                    3. Analyze the answer for any mention of these topics, keywords, or concepts in <thought> tag
                    4. Consider any potential indirect connections or implications that might link the reasoning and the answer. Deductive reasoning need to be valid and sound in <justification> tag
                    5. Evaluate the overall logic and soundness in both the reasoning + context that lead to the answer in <logical> tag

                    As you go through this process, please use a chain of thought approach. Write out your reasoning for each step inside <thought> tags.

                    After your analysis, provide a boolean judgment on whether the reasoning and the answer are logical or relevant to each other. Use "true" if they are logical or relevant, and "false" if they are not.

                    Before giving your final judgment, provide a justification for your decision. Explain the key factors that led to your conclusion.

                    Please ensure your analysis is thorough, impartial, and based on the content provided.
                """,
            },
            {
                "role": "user",
                "content": """
                    Here is the question + context:

                    <reasoning>
                    {{reasoning}}
                    </reasoning>

                    <context>
                    {{context}}
                    </context>
                    
                    Here is the text:
                    <answer>
                    {{answer}}
                    </answer>
                """,
            },
        ],
        response_model=Judgment, #(reasoning=reasoning, context=context, answer=answer)
        context={"reasoning": reasoning, "context": context, "answer": answer},
    )