from trilogy import Environment, Dialects
from trilogy.constants import DEFAULT_NAMESPACE
from trilogy.authoring import (
    Concept,
    ConceptRef,
    SelectStatement,
)
from trilogy.core.models.datasource import Datasource
from trilogy.core.statements.execute import ProcessedShowStatement
from trilogy.core.processing.concept_strategies_v3 import search_concepts, History
from trilogy.executor import Executor
from trilogy.parser import parse_text
from trilogy.core.internal import INTERNAL_NAMESPACE
from trilogy.core.models.build import BuildConcept
from pathlib import Path
from json import loads

example_path = Path(__file__).parent.parent / "examples"


def safe_address(input: Concept | ConceptRef):
    if input.namespace == DEFAULT_NAMESPACE:
        return input.name
    return input.address


def validate_dataset(
    dataset: Datasource, environment: Environment, executor: Executor, dry_run_client
):
    validation_query = (
        "SELECT\n "
        + ",\n\t".join([safe_address(x) for x in dataset.concepts])
        + " LIMIT 0;"
    )
    validate_query(validation_query, environment, executor, dry_run_client)


def validate_query(
    validation_query: str, environment: Environment, executor: Executor, dry_run_client
):
    environment.cte_name_map = {}
    try:
        _, parsed = parse_text(validation_query, environment)
        processed: list[SelectStatement] = [
            x for x in parsed if isinstance(x, SelectStatement)
        ]
        sql = executor.generator.generate_queries(environment, processed)
    except Exception as e:
        print("Failing Validation Query Is")
        print(validation_query)
        raise e
    for statement in sql:
        if isinstance(statement, ProcessedShowStatement):
            continue
        compiled_sql = ""
        # Start the query, passing in the extra configuration.
        try:
            # for UI execution, cap the limit
            compiled_sql = executor.generator.compile_statement(statement)

            # TODO: implement this
            # rs = executor.engine.dry_run(compiled_sql)
            if executor.dialect == Dialects.BIGQUERY:
                # use a dry run to save costs
                # TODO: move this into executor so we don't need this import
                from google.cloud.bigquery import QueryJobConfig

                job_config = QueryJobConfig(dry_run=True, use_query_cache=False)
                query_job = dry_run_client.query(
                    compiled_sql, job_config=job_config
                )  # Make an API request.
                print(
                    "This query will process {} bytes.".format(
                        query_job.total_bytes_processed
                    )
                )
            elif executor.dialect == Dialects.DUCK_DB:
                # use a dry run to save costs
                query_job = executor.execute_raw_sql(compiled_sql)
            elif executor.dialect == Dialects.SNOWFLAKE:
                # dry run with fake snow
                query_job = executor.execute_raw_sql(compiled_sql)
            else:
                raise NotImplementedError(
                    f"Validation not implemented for {executor.dialect}"
                )
        except Exception as e:
            print("Failed validation on:")
            print(validation_query)
            print(compiled_sql)
            raise e


def validate_datasource_grain(datasource):
    # TODO: check that count(*) at datasource grain = 1
    # CON: this requires running an real query
    pass


def get_example_queries(key: str) -> list[str]:
    type, name = key.split(".")
    example_dir = example_path / type / name
    queries = example_dir.glob("*.preql")
    final = []
    for query in queries:
        with open(query, "r") as f:
            final.append(f.read())
    dashboards = example_dir.glob("*.json")
    for dashboard in dashboards:
        with open(dashboard, "r") as f:
            content: dict = loads(f.read())
            imports = content.get("imports", [])
            imp_prefix = " ".join([f"import {x['name']};" for x in imports])
            for k, v in content["gridItems"].items():
                content = v.get("content", None)
                if content and isinstance(content, dict) and content.get("query", None):
                    final.append(imp_prefix + " " + content["query"])
                elif v.get("type") in ["chart", "table"] and v.get("content", None):
                    final.append(imp_prefix + " " + v["content"])

    return final

def validate_concept(concept: BuildConcept, history: History, env, graph):
    if concept.namespace == INTERNAL_NAMESPACE or INTERNAL_NAMESPACE in concept.address:
        return
    search_concepts([concept], history=history, environment=env, depth=0, g=graph)


def validate_model(key: str, model: Environment, executor: Executor, dry_run_client):
    if executor.dialect == Dialects.DUCK_DB:
        return executor.validate_environment()
    # for dataset in model.datasources.values():
    #     validate_dataset(dataset, model, executor, dry_run_client)
    # history = History(base_environment=model)
    # factory = Factory(model)
    # build_model: BuildEnvironment = factory.build(model)
    # graph = generate_graph(build_model)
    # for concept in build_model.concepts.values():
    #     validate_concept(concept, history, build_model, graph)

    for example in get_example_queries(key):

        validate_query(example, model, executor, dry_run_client)
