from __future__ import annotations
import json
import textwrap
import random
import string
import time
from snowflake.snowpark.types import StringType, IntegerType, StructField, StructType, FloatType, MapType, ArrayType, BooleanType, BinaryType, DateType, TimestampType

from snowflake.snowpark.exceptions import SnowparkSQLException

class RAIException(Exception):
    pass

def random_string():
    chars = string.ascii_letters + string.digits
    return ''.join(random.choice(chars) for _ in range(32))

def run_sql(session, query, params=None, engine=""):
    try:
        return session.sql(query, params).collect()
    except SnowparkSQLException as e:
        if "engine not found" in e.message:
            import sys
            def exception_handler(exception_type, exception, traceback):
                print(f"{exception_type.__name__}: {exception}")
            sys.excepthook = exception_handler
            raise RAIException(f"RelationalAI engine not found. Please create an engine called `{engine}` using `___RAI_APP___.api.create_engine`.") from None
        else:
            raise e

def get_models(session, database, engine):
    APP_NAME = {{ APP_NAME }}
    tmp_name = f"tmp_{random_string()}"
    query = textwrap.dedent(f"""
        call {APP_NAME}.api.exec_into(
            '{database}',
            '{engine}',
            'def pairs(name, model): rel(:catalog, :model, name, model) and not starts_with(name, "rel/") and not starts_with(name, "pkg/std")
             def Export_Relation(:key, i, key): exists( (value) | sort(pairs, i, key, value) )
             def Export_Relation(:value, i, value): exists( (key) | sort(pairs, i, key, value) )',
            '{tmp_name}',
            true
        );
    """)
    try:
        run_sql(session, query, engine=engine)
    except Exception as e:
        # this means that there are no models other than the ones in rel/
        return
    result = run_sql(session, f"select key, value from ___RAI_APP___.results.{tmp_name};", engine=engine)
    run_sql(session, f"call {APP_NAME}.api.drop_result_table('{tmp_name}');", engine=engine)
    for row in result:
        if row["KEY"] == "catalog":
            continue
        yield row["KEY"], row["VALUE"]

def get_engine(session, passed_engine):
    if passed_engine:
        return passed_engine
    try:
        APP_NAME = {{ APP_NAME }}
        query = f"select * from {APP_NAME}.api.engines where created_by=current_user() order by created_on DESC;"
        row = run_sql(session, query)[0]
        return row["NAME"]
    except:
        return {{ engine }}

def escape(code):
    return (
        code
        .replace("\\", "\\\\\\\\")
        .replace("'", "\\\'")
        .replace('"', '\\\\"')
        .replace("\n", "\\n")
    )

def get_installation_code(models):
    lines = []
    for (name, code) in models:
        name = escape(name)
        lines.append(textwrap.dedent(f"""
            def delete[:rel, :catalog, :model, "{name}"]: rel[:catalog, :model, "{name}"]
            def insert[:rel, :catalog, :model, "{name}"]: raw\"\"\"\"\"\"\"{code}\"\"\"\"\"\"\"
        """))
    rel_code = "\\n\\n".join(lines)
    return rel_code

    # @NOTE: `overhead_rate` should fall between 0.05 and 0.5 depending on how time sensitive / expensive the operation in question is.
def poll_with_specified_overhead(
    f,
    overhead_rate: float, # This is the percentage of the time we've already waited before we'll poll again.
    start_time: float | None = None,
    timeout: int | None = None,
    max_tries: int | None = None,
    max_delay: float = 120,
    min_delay: float = 0.1
):
    if overhead_rate < 0:
        raise ValueError("overhead_rate must be non-negative")

    if start_time is None:
        start_time = time.time()

    tries = 0
    max_time = time.time() + timeout if timeout else None

    while True:
        if f():
            break

        current_time = time.time()

        if max_tries is not None and tries >= max_tries:
            raise Exception(f'max tries {max_tries} exhausted')

        if max_time is not None and current_time >= max_time:
            raise Exception(f'timed out after {timeout} seconds')

        duration = (current_time - start_time) * overhead_rate
        duration = max(min(duration, max_delay), min_delay, 0)

        time.sleep(duration)
        tries += 1

def check_ready(session, sql_string, database, engine):
    results = run_sql(session, sql_string, [database, engine])
    # Extract the JSON string from the `USE_INDEX` field
    use_index_json_str = results[0]["USE_INDEX"]

    # Convert the JSON string to lowercase
    use_index_json_str = use_index_json_str.lower()

    # Parse the JSON string into a Python dictionary
    use_index_data = json.loads(use_index_json_str)
    ready = use_index_data.get("ready", False)
    errors = use_index_data.get("errors", [])

    if ready:
        return True
    elif errors:
        error_strs = []
        for error in errors:
            if error.get("type") == "data":
                error_strs.append(f"{error.get('message')}, source: {error.get('source')}")
            elif error.get("type") == "engine":
                error_strs.append(f"{error.get('message')}")
        raise Exception("\n".join(error_strs))


def handle(session{{py_inputs}}, passed_engine=""):
    {{clean_inputs}}
    engine = get_engine(session, passed_engine)
    proc_database = {{ proc_database }}
    APP_NAME = {{ APP_NAME }}
    database = {{ database }}
    sql_out_names = "{{ sql_out_names }}"
    models = list(get_models(session, proc_database, engine))
    installation_code = get_installation_code(models)
    rel_code = installation_code + "\n\n" + {{ rel_code }}
    cloned_database = f"{database[:30]}_{random_string()}"
    temp_table = f"temp_{cloned_database}"

    # don't clone the database, instead, call use_index with the given sources
    sources = "{{source_references}}"
    sql_string = f"CALL {APP_NAME}.api.use_index([{sources}], {{'model': ?, 'engine': ?}});"
    # poll this ntil it's ready
    poll_with_specified_overhead(lambda: check_ready(session, sql_string, database, engine), overhead_rate=0.1, max_delay=1)

    try:
        run_sql(session, f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?);", [database, engine, rel_code, cloned_database, False], engine=engine)
        run_sql(session, f"create temporary table {APP_NAME}.results.{temp_table} as select {sql_out_names} from {APP_NAME}.results.{cloned_database};", [], engine=engine)
        run_sql(session, f"call {APP_NAME}.api.drop_result_table(?)", [cloned_database], engine=engine)
        return session.table(f"{APP_NAME}.results.{temp_table}")
    except Exception as e:
        msg = str(e).lower()
        if f"No columns returned".lower() in msg or "Columns of results could not be determined".lower() in msg:
            return session.createDataFrame([], StructType([{{ py_outs }}]))
        raise e
    finally:
        # run_sql(session, f"call {APP_NAME}.api.delete_database('{cloned_database}');")
        pass