import os
from jinja2 import Environment, FileSystemLoader
from typing import List, Dict, Any

# Setup Jinja2 environment
prompt_dir = os.path.join(os.path.dirname(__file__), "prompts")
prompt_env = Environment(loader=FileSystemLoader(prompt_dir))


async def render_evaluate_prompt(evaluate_data: List[Any], schema: Any,
                                 baseline: str = "TIME",
                                 memory_verify: bool = False, quality_assessment: bool = False,
                                 statement_databasets=None, language_type: str = "zh") -> str:
    """
    Renders the evaluate prompt using the evaluate_optimized.jinja2 template.

    Args:
        evaluate_data: The data to evaluate
        schema: The Pydantic model class or JSON schema to use for the output.
        baseline: The baseline type for conflict detection (TIME/FACT/TIME-FACT)
        memory_verify: Whether to enable memory verification for privacy detection

    Returns:
        Rendered prompt content as string
    """
    if statement_databasets is None:
        statement_databasets = []
    template = prompt_env.get_template("evaluate.jinja2")

    # Convert Pydantic model to JSON schema if needed
    if hasattr(schema, 'model_json_schema'):
        json_schema = schema.model_json_schema()
    elif hasattr(schema, 'schema'):
        json_schema = schema.schema()
    else:
        json_schema = schema

    rendered_prompt = template.render(
        evaluate_data=evaluate_data,
        json_schema=json_schema,
        baseline=baseline,
        memory_verify=memory_verify,
        quality_assessment=quality_assessment,
        statement_databasets=statement_databasets,
        language_type=language_type
    )
    return rendered_prompt


async def render_reflexion_prompt(data: Dict[str, Any], schema: Any, baseline: str, memory_verify: bool = False,
                                  statement_databasets=None, language_type: str = "zh") -> str:
    """
    Renders the reflexion prompt using the reflexion_optimized.jinja2 template.

    Args:
        data: The data to reflex on.
        schema: The Pydantic model class or JSON schema to use for the output.
        baseline: The baseline type for conflict resolution.

    Returns:
        Rendered prompt content as a string.
    """
    if statement_databasets is None:
        statement_databasets = []
    template = prompt_env.get_template("reflexion.jinja2")

    # Convert Pydantic model to JSON schema if needed
    if hasattr(schema, 'model_json_schema'):
        json_schema = schema.model_json_schema()
    elif hasattr(schema, 'schema'):
        json_schema = schema.schema()
    else:
        json_schema = schema

    rendered_prompt = template.render(data=data, json_schema=json_schema,
                                      baseline=baseline, memory_verify=memory_verify,
                                      statement_databasets=statement_databasets, language_type=language_type)

    return rendered_prompt
