import dspy
from dspy.functional import TypedPredictor
import logging
from typing import Optional

from app.core.sql_sample_gen import SQLSample

logger = logging.getLogger(__name__)


class SampleExtract(dspy.Signature):
    """Your task is to extract a fully executable SQL example from the provided text that can be used to validate the correctness of the answer in a TiDB Serverless environment. The extracted SQL example should be directly relevant to the user's question and demonstrate the key points of the answer.

    Follow the steps below to ensure the SQL example is actionable and error-free:

    ### Steps to Follow:

    1. **Extract Relevant SQL Example:**
      - Identify and extract the SQL section from the provided content.
      - Ensure the SQL example directly addresses the user's question and the provided answer. It should be a logical sequence of SQL statements that demonstrates the answer's key points.

    2. **Ensure Executability:**
      - Verify that the SQL example can be executed in a TiDB Serverless instance without modifications.
      - Avoid using SQL features not supported in a single TiDB Serverless instance, such as operations requiring multiple clusters, bidirectional replication or complex setups.

    3. **Adjust for TiDB Serverless:**
      - If the SQL contains multiple statements in one line, add `SET SESSION tidb_multi_statement_mode='ON'` before these statements (e.g., 'SET SESSION tidb_enforce_mpp = 1; SELECT ...').
      - Skip unsupported SQL statements like `SPLIT REGION` operations, `tidb_general_log` system variables, `PLACEMENT POLICY` ruls and `mysql.tidb` system tables.

    4. **Setup SQL Execution Environment:**
      - Assume a clean `db` without any pre-existing tables. And avoid assuming any existing `db` names, you don't know them at all.
      - Include steps to create necessary tables before executing dependent statements.
      - Remove database placeholders like `your_database_name`, if possible.
      - eplace placeholders with appropriate values or skip those statements if replacement isn't possible. (e.g., `{your_table_name}`, `{your_column_name}`, `{actual_value}`)

    5. **Error Prevention and Adjustments:**
      - Start with SQL commands that set up the required table structure.
      - Use `DROP TABLE IF EXISTS` before `CREATE TABLE` to avoid conflicts.
      - Adjust any syntax that may not be supported in TiDB Serverless with compatible alternatives.
      - Ensure each SQL statement is self-contained and executable independently.

    6. **Self-Assessment:**
      - Verify the SQL example is relevant to the user's question and answer context.
      - Follow the adjustment rules to make necessary modifications.
      - provide a logically coherent SQL example where each step builds upon the previous one. Pay close attention to the dependencies between the steps to ensure that no dependency restrictions are violated.
    """

    QA_content: str = dspy.InputField(
        desc="The user's query that requires a step-by-step example to be generated.",
        prefix="QA Content:"
    )
    sample: SQLSample = dspy.OutputField(
        desc="Step-by-step example to execute the SQL query in TiDB Serverless.",
        prefix='Here is the executable SQL example for TiDB Serverless:'
    )


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

    def forward(self, QA_content: str):
        with dspy.settings.context(lm=self.dspy_lm):
            return self.prog(
                QA_content=QA_content,
                config={
                    "response_format": {"type": "json_object"}
                },
            )


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

    def gen(self, QA_content: str) -> SQLSample:
        return self.prog(QA_content).sample