import json
import logging
from pathlib import Path
import re

from lumaCLI.metadata.models.bi import (
    Dashboard,
    DashboardManifest,
    DashboardSchemaMetadata,
    DataModel,
)
from lumaCLI.metadata.sources.powerbi.models import DataflowDetails, WorkspaceInfo


logger = logging.getLogger("dlt")


def transform(
    workspace_info: WorkspaceInfo, dataflows_info: list[DataflowDetails]
) -> DashboardManifest:
    # Extract tables from the Power BI metadata.
    tables = extract_tables(workspace_info, dataflows_info)
    reports = extract_reports(workspace_info, tables=tables)

    return DashboardManifest(
        metadata=DashboardSchemaMetadata(schema="dashboard", version=1),
        payload=reports,
    )


def extract_tables(
    workspace_info: WorkspaceInfo, dataflows_info: list[DataflowDetails]
) -> list[dict]:
    tables = []
    # Each dataset table can only have one database table as a source.
    for workspace in workspace_info.workspaces:
        for dataset in workspace.datasets:
            for dataset_table in dataset.tables:
                if dataset_table.source is None:
                    continue

                # Extract the underlying database table.
                source_expression = dataset_table.source[0].expression

                # There are two ways PowerBI can reference Dataflows:
                # 1) The table expression directly walks the PowerPlatform.Dataflows
                #    tree and contains the dataflowId.
                # 2) The table expression references a named dataset source which in
                #    turn references a dataflow entity. Those expressions do NOT
                #    contain the literal 'PowerPlatform.Dataflows' but do contain
                #    an entity reference like: Source{[entity="co_customer"]}[Data]
                # Detect both cases and use the dataflow extractor when appropriate.
                if (
                    "PowerPlatform.Dataflows" in source_expression
                    or 'entity="' in source_expression
                ):
                    table_database_table = _extract_dataflow_table_from_expression(
                        source_expression, dataflows_info
                    )
                else:
                    table_database_table = _extract_table_from_expression(
                        source_expression
                    )

                if not table_database_table:
                    continue

                database_table_name = table_database_table["name"]
                tables.append({
                    "dataset_id": dataset.id,
                    "dataset_table_name": dataset_table.name,
                    "database_table_name": database_table_name,
                    "database_table_schema": table_database_table.get("schema"),
                    "database_table_database": table_database_table.get("database"),
                    "columns": [
                        {"name": column.name, "data_type": column.dataType}
                        for column in dataset_table.columns
                    ],
                    "tags": table_database_table.get("tags", []),
                })
    return tables


def extract_dataflows(dataset):
    """Extract dataset-level Dataflows."""


def extract_reports(
    workspace_info: WorkspaceInfo, tables: list[dict]
) -> list[Dashboard]:
    reports = []
    for workspace in workspace_info.workspaces:
        for report in workspace.reports:
            # We're not interested in PowerBI Apps. Not sure why they're included
            #  - either way, the original report the app is based on is already included
            # in the response.
            if report.name.startswith("[App]"):
                continue

            report_filtered = {}
            report_id = report.id
            report_filtered["external_id"] = report_id
            report_filtered["url"] = (
                "https://app.powerbi.com/groups/" + workspace.id
                or "" + "/reports/" + report_id
            )
            report_filtered["type"] = "powerbi"
            report_filtered["name"] = report.name
            report_filtered["workspace"] = workspace.name
            report_filtered["created_at"] = report.createdDateTime
            report_filtered["modified_at"] = report.modifiedDateTime
            report_filtered["owners"] = [
                {
                    "user_id": user.graphId,
                    "username": user.identifier,
                    "name": user.displayName,
                }
                for user in report.users
                if user.reportUserAccessRight == "Owner"
            ]

            report_tables = [
                {
                    "name": table["database_table_name"],
                    "schema": table["database_table_schema"],
                    "database": table["database_table_database"],
                    "columns": table["columns"],
                    "tags": table["tags"],
                }
                for table in tables
                if table["dataset_id"] == report.datasetId
            ]
            report_filtered["parent_models"] = report_tables

            reports.append(report_filtered)

    return reports


def _extract_table_from_expression(expression: str) -> DataModel | None:
    """Extract schema and table name from expression."""
    # Check if this is a NativeQuery - we don't extract from those
    if "Value.NativeQuery" in expression:
        return None

    # Get database name.
    database_name_expr = re.search(
        r'AmazonRedshift\.Database\s*\(\s*".*?"\s*,\s*"([^" ]+)"\s*\)',
        expression,
        re.IGNORECASE,
    )
    if not database_name_expr:
        return None
    database_name = database_name_expr.group(1).strip()

    # Find source variable assigned to AmazonRedshift.Database
    source_pattern = r"(\w+)\s*=\s*AmazonRedshift\.Database\s*\([^)]+\)"
    source_match = re.search(source_pattern, expression, re.IGNORECASE)

    if not source_match:
        return None

    source_var_name = source_match.group(1).strip()

    # Get schema names.
    schema_pattern = (
        rf'(\w+)\s*=\s*{re.escape(source_var_name)}\s*{{\[Name="([^"]+)"\]}}\[Data\]'
    )
    schema_match = re.findall(schema_pattern, expression, re.IGNORECASE)

    if not schema_match:
        return None

    schema_var_name = schema_match[0][0].strip()  # The variable name given by user.
    schema_name = schema_match[0][1].strip()  # The actual database schema name.

    # Get table metadata.
    table_pattern = (
        r"\w+\s*=\s*" + re.escape(schema_var_name) + r'{\[Name="([^" ]+)"\]}\[Data\]'
    )
    table_match = re.findall(table_pattern, expression, re.IGNORECASE)

    if not table_match:
        return None

    return {
        "name": table_match[0].strip(),
        "schema": schema_name.strip(),
        "database": database_name,
        "tags": [{"source_system": "AmazonRedshift:Direct"}],
    }


def _normalize_shared_name(raw: str) -> str:
    """Strip M quoting, eg. #"Name" or "Name" -> Name."""
    raw = raw.strip()
    raw = raw.rstrip(",")
    if raw.startswith('#"') and raw.endswith('"'):
        return raw[2:-1]
    if raw.startswith('"') and raw.endswith('"'):
        return raw[1:-1]
    return raw


def _extract_entity_block(
    entity: str, document: str, _seen: set | None = None
) -> str | None:
    """Extract the `let ... in ...;` block for a shared entity.

    Handles names like: shared MyEntity, shared "My Entity", shared #"01_Events". Also
    resolves simple 'Source = SomeOtherShared' references by inlining the referenced
    shared block.
    """
    if _seen is None:
        _seen = set()
    # Avoid infinite recursion.
    if entity in _seen:
        return None
    _seen.add(entity)

    # Match shared <name> variants: simple, quoted "name", or #"<name>".
    pattern = rf'shared\s+(?:#"?{re.escape(entity)}"?|"{re.escape(entity)}"|{re.escape(entity)})\s*=\s*let\b(.*?)\bin\s+[^\n;]+;'
    m = re.search(pattern, document, re.DOTALL | re.IGNORECASE)
    if not m:
        return None
    block = m.group(1)

    # Find a 'Source = <rhs>' assignment in the block (first occurrence).
    # Capture whole rhs up to end-of-line (allow commas inside parentheses like
    # Table.Combine({a,b})).
    m_src = re.search(r"(?m)^\s*Source\s*=\s*(?P<rhs>[^\n]+)", block)
    if not m_src:
        return block

    rhs = m_src.group("rhs").strip()

    # If rhs looks like a direct data source, nothing to resolve.
    if re.search(
        r"\b(AmazonRedshift\.Database|Value\.NativeQuery|PowerPlatform\.Dataflows|Csv\.Contents|Sql\.Database)\b",
        rhs,
    ):
        return block

    # Handle cases where the Source is a Table.Combine(...) of several shared tables
    # e.g. Source = Table.Combine({CurrentYear, OneYear, TwoYears}).
    m_tc = re.search(r"Table\.Combine\s*\(\s*\{(?P<items>[^\}]+)\}\s*\)", rhs)
    if m_tc:
        items = m_tc.group("items")
        # Split on commas and normalize each referenced shared name.
        parts = [p.strip() for p in items.split(",") if p.strip()]
        parent_blocks = []
        for p in parts:
            ref = _normalize_shared_name(p)
            if not ref or ref.lower() == entity.lower():
                continue
            pb = _extract_entity_block(ref, document, _seen=_seen)
            if pb:
                parent_blocks.append(pb.rstrip())

        if parent_blocks:
            # Remove the Source = Table.Combine(...) line from current block.
            block_without_source = re.sub(
                r"(?m)^\s*Source\s*=\s*Table\.Combine\s*\([^\n]+\),?\s*\n?",
                "",
                block,
                count=1,
            )
            # Prepend all parent blocks so Database(...) will be visible downstream.
            return "\n".join(parent_blocks) + "\n" + block_without_source.lstrip()

    # Normalize referenced name (remove #, quotes).
    referenced = _normalize_shared_name(rhs)
    # If it references itself or is empty, bail out.
    if not referenced or referenced.lower() == entity.lower():
        return block

    # Try to extract referenced shared block and inline/merge.
    parent_block = _extract_entity_block(referenced, document, _seen=_seen)
    if not parent_block:
        return block

    # Remove the Source = <ref> line from current block to avoid duplicate Source lines.
    block_without_source = re.sub(
        r"(?m)^\s*Source\s*=\s*[^\n,;]+,\s*\n?", "", block, count=1
    )

    # Prepend the parent's let-body so the resulting block contains the real Source
    # definition.
    return parent_block.rstrip() + "\n" + block_without_source.lstrip()


def _get_entity_columns(entity_name: str, dataflow: DataflowDetails) -> list[str]:
    for entity in dataflow.entities:
        if entity["name"] == entity_name:
            return entity["attributes"]
    return []


def _extract_model_from_dataflow_entity_block(
    entity_block: str,
) -> dict[str, str | list[dict]] | None:
    """Extract database, schema and table from an entity block.

    Strategy:
    - Find the database name from AmazonRedshift.Database(..., "dbname", ...)
    - Find all assignments of the form: <lhs> = <rhs>{[Name = "X"]}[Data]
      (lhs/rhs may be quoted like #"Navigation 1").
    - Prefer the first assignment whose RHS equals the variable assigned to
      the Database(...) call as the schema step (e.g. #"Navigation 1" = Source{[Name="schema"]}[Data]).
    - Use the last assignment's Name as the table (most downstream step).
    """
    db_match = re.search(
        r'AmazonRedshift\.Database\(\s*"[^"]*"\s*,\s*"([^\"]+)"',
        entity_block,
        re.IGNORECASE,
    )
    if not db_match:
        return None
    database_name = db_match.group(1).strip()

    # Detect variable assigned to the Database call (usually 'Source').
    src_var_match = re.search(
        r"(\w+)\s*=\s*AmazonRedshift\.Database\s*\(", entity_block, re.IGNORECASE
    )
    source_var_name = src_var_match.group(1).strip() if src_var_match else "Source"

    # Find all assignments like: <lhs> = <rhs>{[Name = "X"]}[Data].
    assign_pattern = r'(\#?"?[\w\d\s\-\_]+"?)\s*=\s*(\#?"?[\w\d\s\-\_]+"?)\s*\{\s*\[\s*Name\s*=\s*"([^\"]+)"\s*\]\s*\}\s*\[Data\]'
    assigns = re.findall(assign_pattern, entity_block, re.IGNORECASE)
    if not assigns:
        return None

    def normalize(n: str) -> str:
        n = n.strip().rstrip(",")
        if n.startswith('#"') and n.endswith('"'):
            return n[2:-1]
        if n.startswith('"') and n.endswith('"'):
            return n[1:-1]
        return n

    normalized = [
        (normalize(lhs), normalize(rhs), name.strip()) for lhs, rhs, name in assigns
    ]

    # Pick schema: first assignment whose RHS is the DB variable (Source), else first
    # assignment.
    schema_name = None
    for _lhs, rhs, name in normalized:
        if rhs.lower() == source_var_name.lower():
            schema_name = name
            break
    if not schema_name:
        schema_name = normalized[0][2]

    # Pick table: last assignment's name (downstream step).
    table_name = normalized[-1][2]

    if not (table_name and schema_name and database_name):
        logger.debug("Failed to extract model from entity block.")
        logger.debug(entity_block)
        return None

    return {
        "name": table_name,
        "schema": schema_name,
        "database": database_name,
        "tags": [{"source_system": "AmazonRedshift:Dataflow"}],
    }


def _extract_dataflow_table_from_expression(
    expression: str, dataflows_info: list[DataflowDetails]
) -> dict[str, str | list[dict]] | None:
    """Extract schema and table name from dataflow expression.

    We first extract the dataflow ID and the name of the entity within the dataflow.

    Next, we parse the dataflow's source code ("document" field in the dataflow's
    section in dataflows_info) to extract the name of the source database table for the
    entity.
    """
    # 1) Try to extract dataflowId directly from the table expression (most common case)
    dataflow_id_pattern = r'dataflowId="([a-f0-9\-]+)"'
    dataflow_id_match = re.search(dataflow_id_pattern, expression, re.IGNORECASE)

    dataflow_id = None
    dataflow_obj = None

    if dataflow_id_match:
        dataflow_id = dataflow_id_match.group(1).strip()

    # 2) Try to extract entity name from the table expression
    entity_pattern = r'{\[entity="([^"]+)"'
    entity_match = re.search(entity_pattern, expression, re.IGNORECASE)
    if not entity_match:
        return None
    entity_name = entity_match.group(1).strip()

    # If we have a dataflow_id, find that dataflow. Otherwise try to resolve via
    # a named source referenced by this table expression.
    if dataflow_id:
        for dataflow in dataflows_info:
            if dataflow.id == dataflow_id:
                dataflow_obj = dataflow
                break
    else:
        # Look for a source reference in the table expression, e.g.
        # 'Source = SourceNewStrategicDim'
        src_ref_match = re.search(
            r'(?m)^\s*Source\s*=\s*(?P<ref>[A-Za-z0-9_#"@]+)', expression
        )
        src_ref = None
        if src_ref_match:
            src_ref = src_ref_match.group("ref").strip()
            # Normalize the name.
            if src_ref.startswith('"') and src_ref.endswith('"'):
                src_ref = src_ref[1:-1]
            if src_ref.startswith('#"') and src_ref.endswith('"'):
                src_ref = src_ref[2:-1]

    if not dataflow_obj:
        # As a last resort, try to find a dataflow that contains an entity with
        # the given name (some datasets reference a named source which in turn
        # points to a dataflow entity; if we don't have the dataset expressions
        # available we can still try matching by entity name across known
        # dataflows).
        for dataflow in dataflows_info:
            for ent in getattr(dataflow, "entities", []) or []:
                ent_name = (
                    ent.get("name")
                    if isinstance(ent, dict)
                    else getattr(ent, "name", None)
                )
                if ent_name and ent_name == entity_name:
                    dataflow_obj = dataflow
                    break
            if dataflow_obj:
                break
        if not dataflow_obj:
            return None

    # Retrieve the dataflow's M code document.
    document = dataflow_obj.pbi_mashup.document

    # Extract the code block containing the entity's definition.
    entity_block = _extract_entity_block(entity_name, document)

    if not entity_block:
        return None

    # Extract the source table from the code block.
    model = _extract_model_from_dataflow_entity_block(entity_block)

    if not model:
        msg = f"Could not extract model from dataflow entity block for entity '{entity_name}'."
        logger.warning(msg)

    return model


if __name__ == "__main__":
    with Path("powerbi_workspace_info.json").open(encoding="utf-8") as f:
        workspace_info = WorkspaceInfo(**next(iter(json.load(f))))

    with Path("powerbi_dataflows_info.json").open(encoding="utf-8") as f:
        dataflows_info = [DataflowDetails(**item) for item in json.load(f)]

    dashboard_manifest = transform(workspace_info, dataflows_info)

    with Path("powerbi_extracted.json").open("w", encoding="utf-8") as f:
        f.write(dashboard_manifest.model_dump_json(by_alias=True))

    # # print(json.dumps(tables, indent=4))
