import dspy

class CheckedResponder(dspy.Module):
    def __init__(self):
        self.respond = dspy.ChainOfThought(GenerateResponse)

    def forward(self, question):
        result = self.respond(question=question)

        # Hard checks — will retry if these fail
        dspy.Assert(
            len(result.answer) > 0,
            "Must produce an answer"
        )
        dspy.Assert(
            len(result.answer.split()) <= 200,
            "Answer must be under 200 words"
        )

        # Soft checks — hints for improvement
        dspy.Suggest(
            "I don't know" not in result.answer.lower(),
            "Try to provide a substantive answer"
        )
        dspy.Suggest(
            not any(word in result.answer.lower() for word in ["definitely", "absolutely", "100%"]),
            "Avoid overconfident language"
        )

        return result

from typing import Literal
from pydantic import BaseModel, Field

class Response(BaseModel):
    answer: str = Field(min_length=1, max_length=500)
    confidence: float = Field(ge=0.0, le=1.0)
    category: str

class MySignature(dspy.Signature):
    question: str = dspy.InputField()
    response: Response = dspy.OutputField()

import re

class ValidatedExtractor(dspy.Module):
    def __init__(self):
        self.extract = dspy.ChainOfThought(ExtractContact)

    def forward(self, text):
        result = self.extract(text=text)

        # Validate email format
        dspy.Assert(
            re.match(r"[^@]+@[^@]+\.[^@]+", result.email or ""),
            "Email must be a valid email address"
        )

        # Validate phone format
        dspy.Assert(
            len(re.sub(r"\D", "", result.phone or "")) >= 10,
            "Phone must have at least 10 digits"
        )

        return result

class VerifyFacts(dspy.Signature):
    """Check if the answer is supported by the given context."""
    context: list[str] = dspy.InputField(desc="Source documents")
    answer: str = dspy.InputField(desc="Generated answer to verify")
    is_supported: bool = dspy.OutputField(desc="Is the answer fully supported by the context?")
    unsupported_claims: list[str] = dspy.OutputField(desc="Claims not found in context")

class GroundedResponder(dspy.Module):
    def __init__(self):
        self.retrieve = dspy.Retrieve(k=5)
        self.answer = dspy.ChainOfThought(AnswerFromDocs)
        self.verify = dspy.Predict(VerifyFacts)

    def forward(self, question):
        context = self.retrieve(question).passages
        response = self.answer(context=context, question=question)

        # Verify the answer is grounded in sources
        check = self.verify(context=context, answer=response.answer)
        dspy.Assert(
            check.is_supported,
            f"Answer contains unsupported claims: {check.unsupported_claims}. "
            "Rewrite using only information from the context."
        )

        return response

class CrossCheckedAnswer(dspy.Module):
    def __init__(self):
        self.answer_a = dspy.ChainOfThought(AnswerQuestion)
        self.answer_b = dspy.ChainOfThought(AnswerQuestion)
        self.compare = dspy.ChainOfThought(CompareAnswers)

    def forward(self, question):
        a = self.answer_a(question=question)
        b = self.answer_b(question=question)

        comparison = self.compare(
            question=question,
            answer_a=a.answer,
            answer_b=b.answer,
        )

        dspy.Assert(
            comparison.agree,
            "Two independent generations disagree — the answer may be unreliable"
        )

        return a

class CompareAnswers(dspy.Signature):
    """Check if two independently generated answers agree."""
    question: str = dspy.InputField()
    answer_a: str = dspy.InputField()
    answer_b: str = dspy.InputField()
    agree: bool = dspy.OutputField(desc="Do the answers substantially agree?")
    discrepancy: str = dspy.OutputField(desc="What they disagree on, if anything")

BLOCKED_PATTERNS = [
    r"\b(password|secret|api.?key)\b",
    r"\b\d{3}-\d{2}-\d{4}\b",  # SSN pattern
]

class SafeResponder(dspy.Module):
    def __init__(self):
        self.respond = dspy.ChainOfThought(GenerateResponse)

    def forward(self, question):
        result = self.respond(question=question)

        # Check for leaked sensitive data
        for pattern in BLOCKED_PATTERNS:
            dspy.Assert(
                not re.search(pattern, result.answer, re.IGNORECASE),
                f"Response may contain sensitive data (pattern: {pattern})"
            )

        return result

class SafetyCheck(dspy.Signature):
    """Check if the response is safe and appropriate."""
    question: str = dspy.InputField()
    response: str = dspy.InputField()
    is_safe: bool = dspy.OutputField()
    concern: str = dspy.OutputField(desc="Safety concern if not safe, empty if safe")

class SafetyCheckedResponder(dspy.Module):
    def __init__(self):
        self.respond = dspy.ChainOfThought(GenerateResponse)
        self.check = dspy.Predict(SafetyCheck)

    def forward(self, question):
        result = self.respond(question=question)

        safety = self.check(question=question, response=result.answer)
        dspy.Assert(
            safety.is_safe,
            f"Response flagged as unsafe: {safety.concern}. Regenerate."
        )

        return result

class FilteredEnsemble(dspy.Module):
    def __init__(self, num_candidates=5):
        self.generators = [dspy.ChainOfThought(GenerateAnswer) for _ in range(num_candidates)]
        self.judge = dspy.ChainOfThought(RankAnswers)

    def forward(self, question):
        candidates = []
        for gen in self.generators:
            try:
                result = gen(question=question)
                # Only keep candidates that pass basic checks
                if len(result.answer) > 0 and len(result.answer.split()) < 200:
                    candidates.append(result.answer)
            except Exception:
                continue

        dspy.Assert(len(candidates) > 0, "No valid candidates generated")

        return self.judge(question=question, candidates=candidates)

class RankAnswers(dspy.Signature):
    """Pick the best answer from the candidates."""
    question: str = dspy.InputField()
    candidates: list[str] = dspy.InputField()
    best_answer: str = dspy.OutputField()