# security/sql_validator.py
import re

_BANNED_KEYWORDS = {"INSERT", "UPDATE", "DELETE", "DROP", "ALTER", "TRUNCATE", "CREATE", "EXEC", "EXECUTE"}
_TABLE_REF_RE = re.compile(r"\bFROM\s+([\[\]\w\.]+)", re.IGNORECASE)

def validate_sql(
    sql: str,
    allowed_tables: set[str],
    default_row_limit: int | None = 50,
) -> str:
    if not sql or not sql.strip():
        raise ValueError("SQL string is empty.")

    sql = sql.strip()

    # 立即拒绝任何非 SELECT 语句
    if not re.match(r"^\s*SELECT\b", sql, re.IGNORECASE):
        raise ValueError(f"Only SELECT queries are allowed. Got: {sql[:80]!r}")

    upper_sql = sql.upper()
    for keyword in _BANNED_KEYWORDS:
        if re.search(rf"\b{keyword}\b", upper_sql):
            raise ValueError(f"Prohibited keyword detected in query: {keyword}")

    # 隔离并校验目标表 schema
    table_refs = _TABLE_REF_RE.findall(sql)
    if not table_refs:
        raise ValueError("Query contains no explicit FROM clause.")

    expanded_allowed = allowed_tables | {t.split(".")[-1] for t in allowed_tables}
    for ref in table_refs:
        normalized = re.sub(r"[\[\]]", "", ref.upper()).strip()
        if normalized not in expanded_allowed:
            raise ValueError(
                f"Unauthorized table reference: {ref!r}. "
                f"Allowed tables: {', '.join(sorted(allowed_tables))}"
                )

    # 自动追加或限制最大行数，以防 OOM 错误
    if default_row_limit is not None and "LIMIT" not in upper_sql:
        sql = f"{sql} LIMIT {default_row_limit}"

    return sql

# llm/provider_manager.py
import logging
from langchain_anthropic import ChatAnthropic
from langchain_core.language_models import BaseChatModel
from langchain_openai import AzureChatOpenAI
from config import CONFIG

logger = logging.getLogger(__name__)

class ProviderManager:
    def __init__(self) -> None:
        self._use_fallback = False
        self.primary: BaseChatModel = ChatAnthropic(
            anthropic_api_url=CONFIG.AZURE_CLAUDE_ENDPOINT,
            anthropic_api_key=CONFIG.AZURE_CLAUDE_API_KEY,
            model=CONFIG.AZURE_CLAUDE_DEPLOYMENT,
            temperature=0,
            max_retries=0,
        )
        self.fallback: BaseChatModel = AzureChatOpenAI(
            azure_endpoint=CONFIG.AZURE_OPENAI_ENDPOINT,
            api_key=CONFIG.AZURE_OPENAI_API_KEY,
            api_version=CONFIG.AZURE_OPENAI_API_VERSION,
            azure_deployment=CONFIG.AZURE_OPENAI_DEPLOYMENT,
            temperature=0,
        )

    @property
    def active(self) -> BaseChatModel:
        return self.fallback if self._use_fallback else self.primary

    @property
    def is_fallback_active(self) -> bool:
        return self._use_fallback

    def switch_to_fallback(self) -> None:
        if not self._use_fallback:
            logger.warning(
                "LLM provider switching to fallback (Azure OpenAI) for remainder of session"
            )
            self._use_fallback = True

provider_manager = ProviderManager()

# llm/failover.py
import asyncio
import logging
from collections.abc import Callable
from typing import Any
from langchain_core.runnables import Runnable
from config import CONFIG
from llm.provider_manager import provider_manager

logger = logging.getLogger(__name__)

async def invoke_with_failover(
    primary_runnable: Runnable,
    fallback_runnable: Runnable,
    inputs: dict,
    invoke_config: dict | None = None,
) -> Any:
    """
    带超时调用 primary_runnable。超时或发生任何异常时，持久切换到备用供应商，并改用备用路径执行。
    """
    if provider_manager.is_fallback_active:
        return await fallback_runnable.ainvoke(inputs, config=invoke_config)

    try:
        return await asyncio.wait_for(
            primary_runnable.ainvoke(inputs, config=invoke_config),
            timeout=float(CONFIG.LLM_FAILOVER_TIMEOUT),
        )
    except Exception as exc:
        logger.error(
            "Primary LLM call failed (%s: %s) — switching to fallback for session",
            type(exc).__name__,
            exc,
        )
        provider_manager.switch_to_fallback()
        return await fallback_runnable.ainvoke(inputs, config=invoke_config)

async def invoke_agent_with_failover(
    build_fn: Callable[[], Any],
    invoke_inputs: dict,
    invoke_config: dict | None = None,
) -> Any:
    """
    构建并调用带故障切换支持的 LangGraph 智能体。
    build_fn 是一个无参可调用对象，用于在备用 LLM 已激活的状态下重新构建智能体拓扑。
    """
    agent = build_fn()

    if provider_manager.is_fallback_active:
        return await agent.ainvoke(invoke_inputs, config=invoke_config)

    try:
        return await asyncio.wait_for(
            agent.ainvoke(invoke_inputs, config=invoke_config),
            timeout=float(CONFIG.LLM_FAILOVER_TIMEOUT),
        )
    except Exception as exc:
        logger.error(
            "Agent LLM call failed (%s: %s) — switching to fallback for session",
            type(exc).__name__,
            exc,
        )
        provider_manager.switch_to_fallback()
        agent = build_fn()
        return await agent.ainvoke(invoke_inputs, config=invoke_config)

# prompt/dynamic_context.py
from langchain_core.vectorstores import VectorStore

async def build_dynamic_few_shot_context(
    user_input: str,
    example_store: VectorStore,
    k: int = 3,
) -> str:
    """
    动态获取语义上最相似的前 k 条历史样例，用以塑造更合适的语气、结构与推理风格。
    """
    relevant_examples = await example_store.similarity_search(user_input, k=k)

    formatted_blocks = []
    for ex in relevant_examples:
        block = f"<example>\nInput: {ex.metadata['input']}\nOutput: {ex.metadata['output']}\n</example>"
        formatted_blocks.append(block)

    formatted = "\n\n".join(formatted_blocks)
    return f"<examples>\n{formatted}\n</examples>"

# eval_sets/sql_generation_chain/manifest.yaml
workflow: sql_generation_chain
version: v1
blob_path: eval-sets/sql_generation_chain/v1.jsonl
thresholds:
  accuracy: 0.95

# eval_sets/rag_workflow/manifest.yaml
workflow: rag_workflow
version: v1
blob_path: eval-sets/rag_workflow/v1.jsonl
thresholds:
  faithfulness: 0.85
  answer_relevancy: 0.80
  context_recall: 0.75

# testing/eval_methods.py
from typing import Any
import pytest

async def test_sql_example(
    workflow: Any,  # 目标 text-to-sql 生成链
    example: Any,   # 含原始输入与标准期望 SQL 的实体
    db_connection: Any,
) -> bool:
    """
    通过将生成查询得到的结果集，与确定性评估数据集上的已知标准答案做对比，来评估查询质量。
    """
    generated_sql = await workflow.generate(example.input)
    validated_query = validate_sql(generated_sql, allowed_tables={"USERS", "METRICS"})

    # 在静态评估数据库快照上并发执行这两条查询
    expected_rows = await db_connection.fetch_all(example.expected_sql)
    actual_rows = await db_connection.fetch_all(validated_query)

    # 对结果行做一致性校验，忽略语法顺序上的细微差异
    assert set(map(tuple, actual_rows)) == set(map(tuple, expected_rows))

# tests/integration/test_quality_gates.py
import os
import json
import yaml
import pytest
import mlflow
from azure.storage.blob import BlobServiceClient
from src.workflows.infosec_rag import InfosecRagChain
from tests.integration.eval_runner import EvalRunner

MANIFEST_PATH = "eval_sets/infosec_rag/manifest.yaml"

@pytest.fixture(scope="module")
def manifest():
    with open(MANIFEST_PATH) as f:
        return yaml.safe_load(f)

@pytest.fixture(scope="module")
def eval_set(manifest):
    client = BlobServiceClient.from_connection_string(os.environ["AZURE_STORAGE_CONNECTION_STRING"])
    blob = client.get_blob_client(container=os.environ["EVAL_BLOB_CONTAINER"], blob=manifest["blob_path"])
    lines = blob.download_blob().readall().decode().strip().splitlines()
    return [json.loads(line) for line in lines]

@pytest.mark.integration
@pytest.mark.asyncio
async def test_infosec_rag_quality(manifest, eval_set):
    runner = EvalRunner(workflow=InfosecRagChain())

    # 会话跟踪细节原生关联到 Git commit 哈希
    with mlflow.start_run(run_name=f"{manifest['workflow']}_{manifest['version']}"):
        mlflow.set_tags({"workflow": manifest["workflow"], "version": manifest["version"]})
        results = await runner.run(eval_set, metrics=list(manifest["thresholds"].keys()))
        mlflow.log_metrics(results.aggregate())

    # 断言结果始终紧扣质量参数
    thresholds = manifest["thresholds"]
    aggregated_metrics = results.aggregate()
    for metric, threshold in thresholds.items():
        actual = aggregated_metrics[metric]
        assert actual >= threshold, f"{metric} has regressed. Found {actual:.3f}, requires >= {threshold}"