src/text2sql/prompt_parser.py

"""
Prompt parser for QAFD cluster output.

Converts KG query clusters into formatted schema documentation
(CREATE TABLE statements with constraints and sample rows).

Matches the CoFD-M prompt.txt format for text2sql pipelines.
"""

import json
from typing import List, Dict, Any, Optional


def parse_qafd_clusters(
    clusters_data: List[Dict[str, Any]],
    add_sample_rows: bool = True,
    schema_data: Optional[Dict] = None,
    format_type: str = "create_table",
) -> str:
    """
    Parse QAFD clusters into formatted schema documentation.

    Args:
        clusters_data: List of cluster dicts from KG query (return_raw_clusters=True)
        add_sample_rows: Whether to include sample rows from schema
        schema_data: Original DB summary JSON for sample rows and column types
        format_type: "create_table" or "simple"

    Returns:
        Formatted schema string
    """
    if not clusters_data:
        return "No schema information available."

    table_info = extract_tables_from_clusters(clusters_data)

    if not table_info:
        return "No tables found in cluster data."

    if format_type == "create_table":
        return format_as_create_table(table_info, schema_data, add_sample_rows)
    else:
        return format_as_simple(table_info, schema_data, add_sample_rows)


def extract_tables_from_clusters(clusters_data: List[Dict[str, Any]]) -> Dict[str, Dict]:
    """Extract table/column info from cluster entities."""
    table_info = {}

    for cluster in clusters_data:
        entities = cluster.get("entities", cluster.get("nodes", []))

        for entity in entities:
            entity_name = entity.get("entity", "").strip('"')
            entity_type = entity.get("entity_type", entity.get("type", ""))
            description = entity.get("description", "")

            if not entity_name:
                continue

            if entity_type == "complete_table":
                if entity_name not in table_info:
                    table_info[entity_name] = {
                        "columns": {},
                        "description": description,
                        "rank": entity.get("rank", 0),
                    }
                elif description and not table_info[entity_name].get("description"):
                    table_info[entity_name]["description"] = description

            elif entity_type == "column" and "." in entity_name:
                parts = entity_name.split(".", 1)
                if len(parts) == 2:
                    table_name, column_name = parts

                    if table_name not in table_info:
                        table_info[table_name] = {
                            "columns": {},
                            "description": "",
                            "rank": 0,
                        }

                    table_info[table_name]["columns"][column_name] = {
                        "description": description,
                        "rank": entity.get("rank", 0),
                    }

    return table_info


def format_as_simple(
    table_info: Dict[str, Dict],
    schema_data: Optional[Dict] = None,
    add_sample_rows: bool = True,
) -> str:
    """Format as simple bullet-point list."""
    output = []

    for table_name, info in sorted(table_info.items()):
        output.append(f"\n{'=' * 60}")
        output.append(f"Table: {table_name}")

        if info.get("description"):
            output.append(f"Description: {info['description']}")

        if info.get("columns"):
            output.append("\nColumns:")
            for col_name, col_info in sorted(info["columns"].items()):
                desc = col_info.get("description", "")
                output.append(f"  - {col_name}: {desc}")

        if add_sample_rows and schema_data:
            samples = _get_sample_rows_json(table_name, schema_data)
            if samples:
                output.append(samples)

    return "\n".join(output)


def format_as_create_table(
    table_info: Dict[str, Dict],
    schema_data: Optional[Dict] = None,
    add_sample_rows: bool = True,
) -> str:
    """Format as CREATE TABLE statements (CoFD-M prompt.txt format)."""
    output = []

    sorted_tables = sorted(
        table_info.items(),
        key=lambda x: x[1].get("rank", 0),
        reverse=True,
    )

    for table_name, info in sorted_tables:
        create_stmt = _generate_create_table(table_name, info, schema_data)
        output.append(create_stmt)

        if add_sample_rows and schema_data:
            samples = _get_sample_rows_json(table_name, schema_data)
            if samples:
                output.append(samples)

        output.append("")

    return "\n".join(output)


def _generate_create_table(
    table_name: str,
    table_info: Dict,
    schema_data: Optional[Dict] = None,
) -> str:
    """
    Generate CREATE TABLE matching CoFD-M prompt.txt format:

    -- table_name
    CREATE TABLE `table_name` (
      `col` TYPE NOT NULL,
      `col2` TYPE PRIMARY KEY,
      ,FOREIGN KEY (`col`) REFERENCES `other_table`(`other_col`),
      ,CHECK (`col` BETWEEN min AND max)
    );
    """
    rag_columns = table_info.get("columns", {})
    table_desc = table_info.get("description", "")

    col_defs = []
    constraints = []

    if schema_data and "tables" in schema_data:
        schema_table = _find_table_in_schema(table_name, schema_data)

        if schema_table and "columns" in schema_table:
            schema_col_lookup = {
                col.get("name", ""): col for col in schema_table["columns"]
            }

            for col_name in rag_columns:
                schema_col = schema_col_lookup.get(col_name, {})
                col_type = schema_col.get("type", "TEXT")

                col_line = f"  `{col_name}` {col_type}"
                if schema_col.get("is_primary_key"):
                    col_line += " PRIMARY KEY"
                elif schema_col.get("not_null"):
                    col_line += " NOT NULL"

                col_defs.append(col_line)

                # FK constraint
                if schema_col.get("is_foreign_key") and schema_col.get("references_table"):
                    ref_table = schema_col["references_table"]
                    ref_col = schema_col.get("references_column", col_name)
                    constraints.append(
                        f"  ,FOREIGN KEY (`{col_name}`) REFERENCES `{ref_table}`(`{ref_col}`)"
                    )

                # CHECK constraint with min/max range
                col_min = schema_col.get("min")
                col_max = schema_col.get("max")
                if col_min is not None and col_max is not None and col_min != col_max:
                    if isinstance(col_min, str):
                        constraints.append(
                            f"  ,CHECK (`{col_name}` BETWEEN '{col_min}' AND '{col_max}')"
                        )
                    else:
                        constraints.append(
                            f"  ,CHECK (`{col_name}` BETWEEN {col_min} AND {col_max})"
                        )
    else:
        # No schema_data — use RAG info only
        for col_name in rag_columns:
            col_defs.append(f"  `{col_name}` TEXT")

    lines = [f"-- {table_name}"]
    lines.append(f"CREATE TABLE `{table_name}` (")

    if col_defs or constraints:
        all_lines = col_defs + constraints
        lines.append(",\n".join(all_lines))
    else:
        lines.append("  -- No columns found")

    lines.append(");")

    return "\n".join(lines)


def _find_table_in_schema(table_name: str, schema_data: Dict) -> Optional[Dict]:
    """Find table in schema, handling qualified names."""
    tables = schema_data.get("tables", {})

    if table_name in tables:
        return tables[table_name]

    short_name = table_name.split(".")[-1]
    for key, value in tables.items():
        if key.split(".")[-1] == short_name:
            return value

    return None


def _get_sample_rows_json(table_name: str, schema_data: Dict) -> str:
    """
    Get sample rows in JSON format (CoFD-M prompt.txt style).

    /* Sample rows:
    [
      {"col1": val1, "col2": val2},
      {"col1": val3, "col2": val4}
    ]
    */
    """
    table_data = _find_table_in_schema(table_name, schema_data)
    if not table_data:
        return ""

    sample_rows = None
    for key in ["sample_rows", "samples", "example_rows", "rows"]:
        if key in table_data:
            sample_rows = table_data[key]
            break

    if not sample_rows:
        return ""

    rows_as_dicts = []

    # Dict-of-lists format
    if isinstance(sample_rows, dict):
        col_names = list(sample_rows.keys())
        max_rows = max(
            (len(v) for v in sample_rows.values() if isinstance(v, list)),
            default=0,
        )
        for i in range(min(3, max_rows)):
            row = {}
            for col in col_names:
                col_data = sample_rows.get(col, [])
                if isinstance(col_data, list) and i < len(col_data):
                    val = col_data[i]
                    if isinstance(val, str) and len(val) > 80:
                        val = val[:77] + "..."
                    row[col] = val
                else:
                    row[col] = None
            rows_as_dicts.append(row)

    # List-of-dicts format
    elif isinstance(sample_rows, list) and sample_rows:
        for row in sample_rows[:3]:
            if isinstance(row, dict):
                cleaned = {}
                for k, v in row.items():
                    if isinstance(v, str) and len(v) > 80:
                        v = v[:77] + "..."
                    cleaned[k] = v
                rows_as_dicts.append(cleaned)

    if not rows_as_dicts:
        return ""

    json_str = json.dumps(rows_as_dicts, indent=2, ensure_ascii=False, default=str)
    return f"/* Sample rows:\n{json_str}\n*/"