parse.py

"""
Parse Tableau TWB/TWBX workbooks and export clean nested JSON for AI consumption.

This script extracts metadata, fields, calculated fields, worksheets, dashboards,
filters, parameters, and lineage relationships from Tableau TWB/TWBX files, then
exports them in a clean hierarchical format optimized for AI context windows.

By default, produces nested JSON (~81% smaller than flat format).
Use --flat to get the verbose internal format with all IDs and edges.

No external dependencies required.

Usage:
  python parse.py --input workbook.twbx --output nested.json
  python parse.py --input workbook.twb --output nested.json --flat
  python parse.py --version
  python parse.py --help
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import re
import sys
import tempfile
import xml.etree.ElementTree as ET
import zipfile
from pathlib import Path
from typing import Iterable

PARSER_VERSION = "1.2.0"
FIELD_TOKEN_PATTERN = re.compile(r"\[([^\]]+)\]")

AGGREGATION_FUNCTIONS = {
    "countd",
    "count",
    "sum",
    "avg",
    "min",
    "max",
    "median",
    "var",
    "variance",
    "stdev",
    "stdevp",
    "counta",
    "countrows",
    "countblank",
    "attr",
    "fixed",
}
CONDITIONAL_KEYWORDS = {"if", "iif", "case", "when", "then", "else", "end"}


def _detect_aggregations(formula: str) -> dict[str, list[str]]:
    result: dict[str, list[str]] = {}
    if not formula:
        return result

    for agg in AGGREGATION_FUNCTIONS:
        pattern = rf"\b{agg}\s*\(\s*\[([^\]]+)\]"
        matches = re.findall(pattern, formula, re.IGNORECASE)
        if matches:
            result[agg] = [_clean_identifier(m) for m in matches]

    return result


def _detect_conditional_usage(formula: str) -> bool:
    if not formula:
        return False
    formula_lower = formula.lower()
    return any(kw in formula_lower for kw in CONDITIONAL_KEYWORDS)


def _extract_column_dependencies(formula: str) -> dict[str, object]:
    tokens = _extract_field_tokens(formula)
    aggregations = _detect_aggregations(formula)
    has_conditional = _detect_conditional_usage(formula)

    all_aggregated: list[str] = []
    for fields in aggregations.values():
        all_aggregated.extend(fields)
    all_aggregated = sorted(set(all_aggregated))

    non_aggregated = sorted({_clean_identifier(t) for t in tokens} - set(all_aggregated))

    return {
        "aggregations": aggregations,
        "has_conditional": has_conditional,
        "aggregated_columns": all_aggregated,
        "non_aggregated_columns": non_aggregated,
    }


def parse_arguments(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Parse Tableau TWB/TWBX and export clean nested JSON for AI consumption."
    )
    parser.add_argument(
        "--version",
        action="version",
        version=f"%(prog)s {PARSER_VERSION}",
    )
    parser.add_argument(
        "--input",
        required=True,
        help="Path to Tableau workbook (.twb or .twbx file).",
    )
    parser.add_argument(
        "--output",
        required=True,
        help="Path to output JSON file.",
    )
    parser.add_argument(
        "--flat",
        action="store_true",
        default=False,
        help="Output verbose flat JSON with all IDs and edges (default: nested format).",
    )
    return parser.parse_args(argv)


def _as_bool(value: str | None) -> bool:
    if value is None:
        return False
    return value.strip().lower() in {"true", "1", "yes"}


def _clean_identifier(value: str | None) -> str:
    if not value:
        return ""
    cleaned = value.strip()
    if cleaned.startswith("[") and cleaned.endswith("]"):
        cleaned = cleaned[1:-1]
    return cleaned.strip().lower()


def _extract_field_tokens(text: str | None) -> list[str]:
    if not text:
        return []
    return [token.strip() for token in FIELD_TOKEN_PATTERN.findall(text) if token.strip()]


def _extract_table_name(relation: ET.Element) -> str:
    # Tableau relation nodes vary by connector. Prefer explicit table/name attrs.
    for attr in ("table", "name", "caption"):
        value = relation.get(attr)
        if value:
            return value
    return "unknown_table"


def extract_twb_from_twbx(input_path: Path) -> Path:
    """
    Extract the first .twb file from a .twbx archive.

    Args:
        input_path: Path to .twbx file

    Returns:
        Path to extracted .twb file (in temp directory)

    Raises:
        ValueError: No .twb file found in archive
        zipfile.BadZipFile: File is not a valid ZIP archive
    """
    try:
        with zipfile.ZipFile(input_path, "r") as zf:
            # Find all .twb files in the archive
            twb_files = sorted([n for n in zf.namelist() if n.lower().endswith(".twb")])

            if not twb_files:
                raise ValueError(f"No .twb file found in archive: {input_path.name}")

            # Extract the first .twb file found
            twb_filename = twb_files[0]
            with tempfile.TemporaryDirectory() as tmpdir:
                zf.extract(twb_filename, tmpdir)
                extracted_path = Path(tmpdir) / twb_filename
                # Read content and return as a temporary file that persists
                # since the caller will use it and then we cleanup
                content = extracted_path.read_bytes()
                # Create a temp file that lives until explicitly cleaned up
                with tempfile.NamedTemporaryFile(suffix=".twb", delete=False) as tmp_file:
                    tmp_file.write(content)
                    return Path(tmp_file.name)
    except zipfile.BadZipFile:
        raise ValueError(f"Invalid ZIP archive: {input_path.name}") from None


def extract_datasources(root: ET.Element) -> list[dict[str, object]]:
    datasources: list[dict[str, object]] = []
    for datasource in root.findall(".//datasource"):
        ds_name = datasource.get("name") or "unknown_datasource"
        ds_caption = datasource.get("caption") or ds_name
        ds_id = f"datasource::{ds_name}"

        tables: list[dict[str, str]] = []
        seen_tables: set[str] = set()
        for relation in datasource.findall(".//relation"):
            relation_type = relation.get("type")
            if relation_type and relation_type.lower() == "join":
                continue
            table_name = _extract_table_name(relation)
            table_id = f"table::{ds_name}::{table_name}"
            if table_id in seen_tables:
                continue
            seen_tables.add(table_id)
            tables.append(
                {
                    "id": table_id,
                    "datasource_id": ds_id,
                    "name": table_name,
                    "relation_type": relation_type or "",
                }
            )

        datasources.append(
            {
                "id": ds_id,
                "name": ds_name,
                "caption": ds_caption,
                "tables": sorted(tables, key=lambda item: item["id"]),
            }
        )

    return sorted(datasources, key=lambda item: item["id"])


def extract_fields(
    root: ET.Element,
) -> tuple[list[dict[str, object]], list[dict[str, object]], list[str]]:
    fields: list[dict[str, object]] = []
    calculated_fields: list[dict[str, object]] = []
    warnings: list[str] = []

    for datasource in root.findall(".//datasource"):
        ds_name = datasource.get("name") or "unknown_datasource"
        ds_id = f"datasource::{ds_name}"
        column_nodes = datasource.findall(".//column")

        for column in column_nodes:
            raw_name = column.get("name") or ""
            if not raw_name:
                continue

            column_id = f"field::{ds_name}::{raw_name}"
            caption = column.get("caption") or raw_name
            datatype = column.get("datatype") or ""
            role = column.get("role") or ""
            hidden = _as_bool(column.get("hidden"))
            field_kind = column.get("type") or "column"

            field_record: dict[str, object] = {
                "id": column_id,
                "datasource_id": ds_id,
                "raw_name": raw_name,
                "name": _clean_identifier(raw_name),
                "caption": caption,
                "datatype": datatype,
                "role": role,
                "hidden": hidden,
                "kind": field_kind,
            }
            fields.append(field_record)

            calc_node = column.find("calculation")
            if calc_node is None:
                continue

            formula = calc_node.get("formula") or ""
            calc_id = f"calc::{ds_name}::{raw_name}"
            tokens = sorted({_clean_identifier(t) for t in _extract_field_tokens(formula) if t})
            if not formula:
                warnings.append(f"Calculated field '{raw_name}' in datasource '{ds_name}' has empty formula.")

            column_deps = _extract_column_dependencies(formula)

            calculated_fields.append(
                {
                    "id": calc_id,
                    "field_id": column_id,
                    "datasource_id": ds_id,
                    "raw_name": raw_name,
                    "name": _clean_identifier(raw_name),
                    "caption": caption,
                    "formula": formula,
                    "upstream_field_tokens": tokens,
                    "aggregations": column_deps.get("aggregations", {}),
                    "has_conditional": column_deps.get("has_conditional", False),
                    "aggregated_columns": column_deps.get("aggregated_columns", []),
                    "non_aggregated_columns": column_deps.get("non_aggregated_columns", []),
                }
            )

    return (
        sorted(fields, key=lambda item: item["id"]),
        sorted(calculated_fields, key=lambda item: item["id"]),
        sorted(set(warnings)),
    )


def extract_worksheets(root: ET.Element) -> list[dict[str, object]]:
    """
    Extract worksheets with visualization details.

    For each worksheet, extract:
    - name
    - referenced fields
    - rows/columns (what fields are on each axis)
    - mark type (bar, line, pie, etc.)
    - filters applied at worksheet level
    - encodings (color, size, shape, text)
    - datasource references
    """
    worksheets: list[dict[str, object]] = []

    for worksheet in root.findall(".//worksheet"):
        sheet_name = worksheet.get("name") or "unknown_worksheet"
        sheet_id = f"worksheet::{sheet_name}"

        # Extract referenced fields (existing logic)
        refs: set[str] = set()
        for node in worksheet.iter():
            for attr_value in node.attrib.values():
                refs.update(_clean_identifier(token) for token in _extract_field_tokens(attr_value))
            column_attr = node.get("column")
            if column_attr:
                refs.add(_clean_identifier(column_attr))
        refs.discard("")

        # Extract visualization details
        table = worksheet.find(".//table")
        rows_fields: list[str] = []
        cols_fields: list[str] = []
        mark_type: str = "Automatic"
        encodings: dict[str, str] = {}
        ds_refs: list[str] = []

        if table is not None:
            # Extract rows/columns
            rows_elem = table.find("rows")
            if rows_elem is not None and rows_elem.text:
                rows_fields = _extract_field_tokens(rows_elem.text)

            cols_elem = table.find("cols")
            if cols_elem is not None and cols_elem.text:
                cols_fields = _extract_field_tokens(cols_elem.text)

            # Extract mark type
            for mark in table.findall(".//mark"):
                mark_class = mark.get("class")
                if mark_class:
                    mark_type = mark_class
                    break

            # Extract encodings (color, size, shape, text)
            for enc in table.findall(".//encodings"):
                for child in list(enc):
                    enc_type = child.tag.split("}")[-1] if "}" in child.tag else child.tag
                    col = child.get("column", "")
                    if col:
                        # Clean up the column reference
                        encodings[enc_type] = col

            # Extract datasource references
            for ds_ref in table.findall(".//datasource"):
                ds_name = ds_ref.get("name")
                if ds_name:
                    ds_refs.append(ds_name)

        # Extract filters
        filters = _extract_worksheet_filters(worksheet)

        worksheets.append(
            {
                "id": sheet_id,
                "name": sheet_name,
                "referenced_fields": sorted(refs),
                "rows": rows_fields,
                "cols": cols_fields,
                "mark_type": mark_type,
                "encodings": encodings,
                "filters": filters,
                "datasources": ds_refs,
            }
        )

    return sorted(worksheets, key=lambda item: item["id"])


def _extract_worksheet_filters(worksheet: ET.Element) -> list[dict[str, object]]:
    """
    Extract filters from a worksheet element.

    Filter types:
    - categorical: dimension filters with specific values
    - relative-date: relative date filters (last N days/months)
    - quantitative: measure filters (> 100, between X and Y)
    - range: date/number range filters
    """
    filters: list[dict[str, object]] = []

    for filter_elem in worksheet.findall(".//filter"):
        filter_class = filter_elem.get("class", "unknown")
        filter_column = filter_elem.get("column", "")

        # Clean up column reference
        column_name = _extract_field_tokens(filter_column)
        column_clean = column_name[0] if column_name else filter_column

        filter_record: dict[str, object] = {
            "type": filter_class,
            "field": column_clean,
            "field_raw": filter_column,
        }

        # Extract filter values/members
        members: list[str] = []
        for groupfilter in filter_elem.findall(".//groupfilter"):
            func = groupfilter.get("function", "")
            member_value = groupfilter.get("member", "")
            level = groupfilter.get("level", "")

            if func == "member" and member_value:
                # Boolean or single value filter
                filter_record["function"] = func
                filter_record["member"] = member_value
            elif func == "level-members":
                # All members filter
                filter_record["function"] = func
                filter_record["level"] = _extract_field_tokens(level)[0] if level else level
            elif func in ("intersection", "union"):
                # Multiple values filter - collect members
                for m in groupfilter.findall(".//member"):
                    if m.text:
                        members.append(m.text.strip('"').strip("'"))
                if members:
                    filter_record["values"] = members
                    filter_record["function"] = "include"

            # Check for context filter
            if filter_elem.get("context") == "true":
                filter_record["context"] = True

        # Handle relative-date filters
        if filter_class == "relative-date":
            period_type = filter_elem.get("period-type-v2", "")
            first_period = filter_elem.get("first-period", "")
            last_period = filter_elem.get("last-period", "")
            if period_type:
                filter_record["period_type"] = period_type
            if first_period and last_period:
                filter_record["range"] = f"{first_period} to {last_period}"
            # Include future/null flags
            if filter_elem.get("include-future") == "true":
                filter_record["include_future"] = True
            if filter_elem.get("include-null") == "true":
                filter_record["include_null"] = True

        # Handle quantitative filters
        if filter_class == "quantitative":
            included = filter_elem.get("included-values", "")
            if included:
                filter_record["included_values"] = included
            # Try to find range bounds
            for range_elem in filter_elem.findall(".//range"):
                min_val = range_elem.get("min")
                max_val = range_elem.get("max")
                if min_val or max_val:
                    filter_record["bounds"] = {"min": min_val, "max": max_val}

        # Add any extracted members to the filter
        if members and "values" not in filter_record:
            filter_record["values"] = members

        filters.append(filter_record)

    return filters


def extract_dashboards(root: ET.Element) -> list[dict[str, object]]:
    """
    Extract dashboards with their zones (worksheet containers).
    """
    dashboards: list[dict[str, object]] = []

    for dashboard in root.findall(".//dashboard"):
        db_name = dashboard.get("name") or "unknown_dashboard"
        db_id = f"dashboard::{db_name}"

        zones: list[dict[str, object]] = []
        for zone in dashboard.findall(".//zone"):
            zone_name = zone.get("name", "")
            zone_id = zone.get("id", "")
            zone_type = zone.get("type", "")

            # Parse position/size if available
            x = zone.get("x")
            y = zone.get("y")
            w = zone.get("w")
            h = zone.get("h")

            zone_record: dict[str, object] = {
                "name": zone_name if zone_name else None,
                "type": zone_type if zone_type else None,
            }
            if zone_id:
                zone_record["id"] = zone_id
            if x and y:
                zone_record["position"] = {"x": int(x), "y": int(y)}
            if w and h:
                zone_record["size"] = {"width": int(w), "height": int(h)}

            # Only include zones with useful info
            if zone_name or zone_type:
                zones.append(zone_record)

        dashboards.append(
            {
                "id": db_id,
                "name": db_name,
                "zones": zones,
            }
        )

    return sorted(dashboards, key=lambda item: item["id"])


def extract_parameters(root: ET.Element) -> list[dict[str, object]]:
    """
    Extract parameters from the workbook.

    Parameters are special fields that users can change interactively.
    """
    parameters: list[dict[str, object]] = []
    seen_params: set[str] = set()

    for datasource in root.findall(".//datasource"):
        # Check if this is a Parameters datasource
        ds_name = datasource.get("name", "")
        if ds_name != "Parameters":
            continue

        for column in datasource.findall(".//column"):
            param_domain = column.get("param-domain-type")
            if not param_domain:
                continue

            raw_name = column.get("name", "")
            caption = column.get("caption", raw_name)
            datatype = column.get("datatype", "")

            # Get the formula/value
            calc_node = column.find("calculation")
            formula = calc_node.get("formula", "") if calc_node is not None else ""

            param_id = f"parameter::{raw_name}"
            if param_id in seen_params:
                continue
            seen_params.add(param_id)

            parameters.append(
                {
                    "id": param_id,
                    "name": _clean_identifier(raw_name),
                    "raw_name": raw_name,
                    "caption": caption,
                    "datatype": datatype,
                    "domain_type": param_domain,
                    "current_value": formula,
                }
            )

    return parameters


def extract_relationships(root: ET.Element) -> list[dict[str, object]]:
    """
    Extract relationships (joins) between tables.
    """
    relationships: list[dict[str, object]] = []

    for relation in root.findall(".//relation"):
        rel_type = relation.get("type", "")
        if rel_type.lower() != "join":
            continue

        rel_name = relation.get("name", "") or relation.get("table", "")

        # Extract join clauses if present
        join_clauses: list[dict[str, str]] = []
        for clause in relation.findall(".//expression"):
            left = clause.get("left", "")
            right = clause.get("right", "")
            if left and right:
                join_clauses.append({"left": left, "right": right})

        relationships.append(
            {
                "type": "join",
                "name": rel_name,
                "clauses": join_clauses,
            }
        )

    return relationships


def _field_index(
    fields: Iterable[dict[str, object]],
    calculated_fields: Iterable[dict[str, object]],
) -> tuple[dict[str, str], dict[str, str]]:
    field_by_token: dict[str, str] = {}
    calc_by_token: dict[str, str] = {}

    for field in fields:
        token = _clean_identifier(str(field.get("raw_name", "")))
        if token and token not in field_by_token:
            field_by_token[token] = str(field["id"])

    for calc in calculated_fields:
        token = _clean_identifier(str(calc.get("raw_name", "")))
        if token and token not in calc_by_token:
            calc_by_token[token] = str(calc["id"])

    return field_by_token, calc_by_token


def build_lineage(
    datasources: list[dict[str, object]],
    fields: list[dict[str, object]],
    calculated_fields: list[dict[str, object]],
    worksheets: list[dict[str, object]],
) -> tuple[list[dict[str, str]], list[str]]:
    edges: set[tuple[str, str, str]] = set()
    warnings: list[str] = []

    fields_by_datasource: dict[str, list[str]] = {}
    for field in fields:
        ds_id = str(field["datasource_id"])
        fields_by_datasource.setdefault(ds_id, []).append(str(field["id"]))

    for datasource in datasources:
        ds_id = str(datasource["id"])
        for table in datasource.get("tables", []):  # type: ignore[not-iterable]
            table_id = str(table["id"])
            edges.add((ds_id, table_id, "datasource_to_table"))
            for field_id in fields_by_datasource.get(ds_id, []):
                edges.add((table_id, field_id, "table_to_field"))

    field_by_token, calc_by_token = _field_index(fields, calculated_fields)
    calc_ids = {str(calc["id"]) for calc in calculated_fields}

    for calc in calculated_fields:
        calc_id = str(calc["id"])
        field_id = str(calc["field_id"])
        edges.add((field_id, calc_id, "field_defines_calculation"))
        for token in calc.get("upstream_field_tokens", []):  # type: ignore[not-iterable]
            token_name = _clean_identifier(str(token))
            upstream_id = calc_by_token.get(token_name) or field_by_token.get(token_name)
            if not upstream_id:
                warnings.append(
                    f"Unresolved upstream token '{token_name}' in calculated field '{calc.get('raw_name', '')}'."
                )
                continue
            relation = "calculation_to_calculation" if upstream_id in calc_ids else "field_to_calculation"
            edges.add((upstream_id, calc_id, relation))

    for worksheet in worksheets:
        sheet_id = str(worksheet["id"])
        for token in worksheet.get("referenced_fields", []):  # type: ignore[not-iterable]
            token_name = _clean_identifier(str(token))
            target_id = calc_by_token.get(token_name) or field_by_token.get(token_name)
            if not target_id:
                continue
            relation = "calculation_to_worksheet" if target_id in calc_ids else "field_to_worksheet"
            edges.add((target_id, sheet_id, relation))

    edge_list = [
        {"from": edge[0], "to": edge[1], "relation_type": edge[2]}
        for edge in sorted(edges, key=lambda item: (item[0], item[2], item[1]))
    ]

    return edge_list, sorted(set(warnings))


def build_output(
    input_path: Path,
    datasources: list[dict[str, object]],
    fields: list[dict[str, object]],
    calculated_fields: list[dict[str, object]],
    worksheets: list[dict[str, object]],
    dashboards: list[dict[str, object]],
    parameters: list[dict[str, object]],
    relationships: list[dict[str, object]],
    edges: list[dict[str, str]],
    warnings: list[str],
    original_filename: str | None = None,
) -> dict[str, object]:
    workbook_name = original_filename if original_filename else input_path.name
    return {
        "metadata": {
            "workbook_file": workbook_name,
            "workbook_path": str(input_path),
            "parsed_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
            "parser_version": PARSER_VERSION,
        },
        "datasources": datasources,
        "fields": fields,
        "calculated_fields": calculated_fields,
        "worksheets": worksheets,
        "dashboards": dashboards,
        "parameters": parameters,
        "relationships": relationships,
        "lineage_edges": edges,
        "warnings": warnings,
    }


def parse_workbook(input_path: Path, original_filename: str | None = None) -> dict[str, object]:
    tree = ET.parse(input_path)
    root = tree.getroot()

    datasources = extract_datasources(root)
    fields, calculated_fields, extraction_warnings = extract_fields(root)
    worksheets = extract_worksheets(root)
    dashboards = extract_dashboards(root)
    parameters = extract_parameters(root)
    relationships = extract_relationships(root)
    lineage_edges, graph_warnings = build_lineage(
        datasources=datasources,
        fields=fields,
        calculated_fields=calculated_fields,
        worksheets=worksheets,
    )

    combined_warnings = sorted(set(extraction_warnings + graph_warnings))
    return build_output(
        input_path=input_path,
        datasources=datasources,
        fields=fields,
        calculated_fields=calculated_fields,
        worksheets=worksheets,
        dashboards=dashboards,
        parameters=parameters,
        relationships=relationships,
        edges=lineage_edges,
        warnings=combined_warnings,
        original_filename=original_filename,
    )


def main(argv: list[str]) -> int:
    args = parse_arguments(argv)
    input_path = Path(args.input)
    output_path = Path(args.output)

    ext = input_path.suffix.lower()
    if ext not in (".twb", ".twbx"):
        print(f"Input must be a .twb or .twbx file, received: {input_path}", file=sys.stderr)
        return 2
    if not input_path.exists():
        print(f"Input file does not exist: {input_path}", file=sys.stderr)
        return 2

    # Store original filename for TWBX to preserve in metadata
    original_filename = input_path.name if ext == ".twbx" else None

    # Handle .twbx (ZIP archive) extraction
    extracted_path: Path | None = None
    try:
        if ext == ".twbx":
            try:
                extracted_path = extract_twb_from_twbx(input_path)
                input_path = extracted_path
            except ValueError as e:
                print(str(e), file=sys.stderr)
                return 2
            except zipfile.BadZipFile:
                print(f"Invalid ZIP archive: {input_path.name}", file=sys.stderr)
                return 2

        # Parse the workbook into internal data structure
        flat_data = parse_workbook(input_path, original_filename=original_filename)

        # Transform to nested format unless --flat is requested
        if args.flat:
            output_data = flat_data
        else:
            # Lazy import to avoid circular dependency concerns
            from export_nested import export_nested

            output_data = export_nested(flat_data)

        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(
            json.dumps(output_data, indent=2, sort_keys=not args.flat, ensure_ascii=False),
            encoding="utf-8",
        )

        # Print summary
        if args.flat:
            print(f"Flat lineage JSON written to: {output_path}")
        else:
            print(f"Nested JSON written to: {output_path}")
            counts = output_data.get("counts", {})
            print(f"  Datasources: {counts.get('datasources', 0)}")
            print(f"  Fields: {counts.get('fields', 0)}")
            print(f"  Calculated fields: {counts.get('calculated_fields', 0)}")
            print(f"  Worksheets: {counts.get('worksheets', 0)}")
            print(f"  Dashboards: {counts.get('dashboards', 0)}")
            print(f"  Parameters: {counts.get('parameters', 0)}")

        return 0
    finally:
        # Cleanup extracted temp file
        if extracted_path and extracted_path.exists():
            try:
                extracted_path.unlink()
            except OSError:
                pass


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))