import dspy

from dspy.functional import TypedPredictor
from pydantic import BaseModel, Field
from typing import Optional

from app.models.validation_job import Step


class VerificationResult(BaseModel):
    excepted: bool = Field(
        description="Whether the SQL query results satisfy the criteria outlined in the expected SQL result explanation."
    )
    message: str = Field(description="Message explaining the verification result")


class SQLVerify(dspy.Signature):
    """Your task is to verify whether the SQL query results meet the expected criteria based on the provided execution history and explanation. Follow these rules to assess the SQL queries accurately:

    1. SQL Execution Failure:
      - If the SQL query is expected to fail (as per the explanation), and it does fail, consider it a success for validation purposes.
    2. SQL Execution Success with Conflicting Results:
      - If the SQL query executes successfully, but the result contradicts the expected outcome or the execution history indicates a failure, consider it a failure.
    3. SQL Execution Success with Unpredictable Results:
      - If the SQL query executes successfully, but its result is not derived from a previous operation in the execution history, focus on the query's executability rather than its specific outcome.
      - Do not assume the existence of data, objects, or operations that are not explicitly created or modified in the execution history.
      - For queries that rely on unpredictable environments (like DDL for TiFlash replicas or existing large datasets), your primary focus should be on whether the SQL query executes without errors.
    4. Special case: If you encounter a commented SQL statement, consider it as meeting the expected criteria regardless of its execution result, such as \
        `sql: -- CREATE INDEX idx_embedding ON vector_data USING HNSW(embedding); result: [["Query OK, 2 rows affected (0.159 sec)"]]`
    """

    context = dspy.InputField(desc="Execution history")
    sql = dspy.InputField(desc="SQL query")
    expected_sql_result_explanation = dspy.InputField(
        desc="Expected SQL result explanation"
    )
    sql_result = dspy.InputField(desc="SQL query results")
    result: VerificationResult = dspy.OutputField(desc="Verification result")


class SQLVerifyModule(dspy.Module):
    def __init__(self, dspy_lm: dspy.LM):
        super().__init__()
        self.dspy_lm = dspy_lm
        self.prog = TypedPredictor(SQLVerify)

    def forward(self, context: str, step: Step, results):
        with dspy.settings.context(lm=self.dspy_lm):
            return self.prog(
                context=context,
                sql=step.sql,
                expected_sql_result_explanation=step.explanation,
                sql_result=results,
            )


class SQLVerifier:
    def __init__(self, dspy_lm: dspy.LM, complied_program_path: Optional[str] = None):
        self.prog = SQLVerifyModule(dspy_lm)
        if complied_program_path is not None:
            self.prog.load(complied_program_path)

    def verify(self, context: str, step: Step, results):
        response = self.prog(context, step, results)
        return response.result