#!/usr/bin/env python3
###############################################################################
# MIT License
#
# Copyright (c) 2025 Advanced Micro Devices, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
###############################################################################

import argparse
import os
import sqlite3
import time

from typing import List, Dict, Iterable, Optional, Callable, Any


def merge_sqlite_dbs(
    sources: Iterable[str],
    dest_path: str,
    on_log: Optional[Callable[[str], None]] = None,
) -> None:
    """
    Merge multiple SQLite databases into a single destination database.

    Parameters
    ----------
    sources : Iterable[str]
        Paths to source databases.
    dest_path : str
        Path to destination database.
    on_log : Optional[Callable[[str], None]]
        Logger function; defaults to None. Pass `print` to generate logs.
    """

    def log(msg: str) -> None:
        if on_log:
            on_log(f"  {msg}")

    sources = list(sources)
    if not sources:
        raise ValueError("No source databases provided")

    # Prepare output directory
    output_dir = os.path.dirname(os.path.abspath(dest_path)) or os.getcwd()
    os.makedirs(output_dir, exist_ok=True)

    # Remove existing file
    if os.path.isfile(dest_path):
        os.remove(dest_path)

    uuids = []
    views = []
    data_views = []
    schema_versions = []

    with sqlite3.connect(str(dest_path)) as conn:
        conn.execute("PRAGMA journal_mode = WAL;")
        conn.execute("PRAGMA synchronous = NORMAL;")
        conn.execute("PRAGMA foreign_keys = OFF;")  # defer FK checks until end

        # One big atomic transaction
        with conn:
            # Attach sources one by one
            for i, src in enumerate(sources, 1):
                alias = f"src{i}"
                conn.execute(f"ATTACH DATABASE ? AS {alias}", (src,))
                print(f"Adding {src}")
                log(f"Attached {src} AS {alias}")

                # UUIDs and schema version
                _uuids = [
                    itr[0]
                    for itr in conn.execute(
                        f"SELECT value FROM {alias}.rocpd_metadata WHERE tag='uuid'",
                    ).fetchall()
                ]
                uuids += [itr for itr in _uuids if itr not in uuids]

                _schema_versions = [
                    itr[0]
                    for itr in conn.execute(
                        f"SELECT value FROM {alias}.rocpd_metadata WHERE tag='schema_version'",
                    ).fetchall()
                ]
                schema_versions += _schema_versions

                # Helper: fetch rows from attached sqlite_master
                def fetch_master(_alias: str, kind: str):
                    cur = conn.execute(
                        f"""
                        SELECT name, sql
                        FROM {_alias}.sqlite_master
                        WHERE type = ? AND name NOT LIKE 'sqlite_%'
                        ORDER BY name
                        """,
                        (kind,),
                    )
                    return cur.fetchall()

                # Track dest tables to detect collisions quickly
                existing_tables = {
                    row[0]
                    for row in conn.execute(
                        "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
                    )
                }

                # 1) Create tables
                for name, create_sql in fetch_master(alias, "table"):
                    if name in existing_tables:
                        raise AssertionError(
                            f"Table name collision for '{name}' from {alias}; "
                            "assumption of globally-unique table names violated."
                        )
                    if not create_sql:
                        continue
                    log(f"Creating table {name}")
                    conn.execute(create_sql)
                    existing_tables.add(name)

                # 2) Copy table data
                tbls = [name for name, _ in fetch_master(alias, "table")]
                print(f"Tables found: {len(tbls)}")
                for name in tbls:
                    log(f"Inserting rows into {name} from {alias}.{name}")
                    rows = conn.execute(f'SELECT * FROM {alias}."{name}"').fetchall()
                    if rows:
                        col_count = len(rows[0])
                        placeholders = ", ".join(["?"] * col_count)
                        conn.executemany(
                            f'INSERT INTO "{name}" VALUES ({placeholders})', rows
                        )

                # 3) Recreate indexes (make idempotent with IF NOT EXISTS)
                def inject_if_not_exists_in_index_sql(sql: str) -> str:
                    # Naive, but works for standard forms produced by sqlite_master
                    # Handles UNIQUE and non-UNIQUE:
                    # "CREATE INDEX name ON ..." or "CREATE UNIQUE INDEX name ON ..."
                    sql_stripped = sql.strip()
                    if sql_stripped.upper().startswith("CREATE UNIQUE INDEX"):
                        return sql_stripped.replace(
                            "CREATE UNIQUE INDEX", "CREATE UNIQUE INDEX IF NOT EXISTS", 1
                        )
                    if sql_stripped.upper().startswith("CREATE INDEX"):
                        return sql_stripped.replace(
                            "CREATE INDEX", "CREATE INDEX IF NOT EXISTS", 1
                        )
                    return sql

                existing_indexes = {
                    row[0]
                    for row in conn.execute(
                        "SELECT name FROM sqlite_master WHERE type='index' AND name NOT LIKE 'sqlite_%'"
                    )
                }
                for name, create_sql in fetch_master(alias, "index"):
                    if not create_sql:
                        continue
                    if name in existing_indexes:
                        log(f"Index {name} exists; skipping or using IF NOT EXISTS")
                    # Try to create with IF NOT EXISTS to avoid collision
                    sql2 = inject_if_not_exists_in_index_sql(create_sql)
                    conn.execute(sql2)
                    existing_indexes.add(name)

                # 4) Recreate triggers (skip on name conflict)
                existing_triggers = {
                    row[0]
                    for row in conn.execute(
                        "SELECT name FROM sqlite_master WHERE type='trigger'"
                    )
                }
                for name, create_sql in fetch_master(alias, "trigger"):
                    if not create_sql:
                        continue
                    if name in existing_triggers:
                        log(f"Trigger {name} exists; skipping")
                        continue
                    log(f"Creating trigger {name}")
                    conn.execute(create_sql)
                    existing_triggers.add(name)

                # 5) Recreate views (skip on name conflict)
                existing_views = {
                    row[0]
                    for row in conn.execute(
                        "SELECT name FROM sqlite_master WHERE type='view'"
                    )
                }
                for name, create_sql in fetch_master(alias, "view"):
                    if not create_sql:
                        continue
                    if name in existing_views:
                        log(f"View {name} exists; skipping")
                        continue
                    # If the view name does not start with "rocpd_", collect it for later recreation
                    if not name.startswith("rocpd_") and not any(
                        name == _name for _name, _ in data_views
                    ):
                        data_views.append((name, create_sql))
                    existing_views.add(name)

                views += [itr for itr in list(existing_views) if itr.startswith("rocpd_")]

                conn.commit()
                conn.execute(f"DETACH DATABASE {alias}")
                log(f"Detached {alias}")

        # Check the schema versions.  Merge only occurs if all the DBs are the same schema version.
        unique_versions = list(set(schema_versions))
        if len(unique_versions) != 1:
            raise RuntimeError(f"Multiple schema versions found: {unique_versions}")

        # Re-enable FKs and run a quick FK check
        conn.execute("PRAGMA foreign_keys = ON;")
        # Optional: enforce integrity
        # try:
        #     conn.execute("PRAGMA quick_check;")
        # except sqlite3.DatabaseError as e:
        #     log(f"SQLite3 quick_check reported an issue: {e}")

        uuids = sorted(list(set(uuids)))  # unique set of uuids
        views = sorted(list(set(views)))  # unique set of views

        # Create UNION views by listing all tables
        existing_tables = {
            row[0]
            for row in conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
            )
        }

        # Then UNION all the tables starting with the view name
        for vitr in views:
            matching_tables = [
                titr for titr in existing_tables if titr.startswith(f"{vitr}_")
            ]
            tables_union = " UNION ALL ".join(
                [f"SELECT * FROM {titr}" for titr in matching_tables]
            )
            conn.execute(f"CREATE VIEW {vitr} AS {tables_union}")
        conn.commit()

        # Now that the rocpd_ views are created, re-create the data-views using all the data
        for _, sql_view in data_views:
            conn.execute(sql_view)
        conn.commit()


#
# Command-line interface functions
#
def add_args(parser):
    """Add arguments for merger."""

    io_options = parser.add_argument_group("I/O options")

    io_options.add_argument(
        "-o",
        "--output-file",
        help="Sets the base output file name",
        default=os.environ.get("ROCPD_OUTPUT_NAME", "merged"),
        type=str,
        required=False,
    )
    io_options.add_argument(
        "-d",
        "--output-path",
        help="Sets the output path where the output files will be saved (default path: `./rocpd-output-data`)",
        default=os.environ.get("ROCPD_OUTPUT_PATH", "./rocpd-output-data"),
        type=str,
        required=False,
    )

    def process_args(input, args):
        valid_args = ["output_file", "output_path"]
        ret = {}
        for itr in valid_args:
            if hasattr(args, itr):
                val = getattr(args, itr)
                if val is not None:
                    ret[itr] = val
        return ret

    return process_args


def execute(inputs: List[str], **kwargs: Dict[str, Any]) -> str:

    start_time = time.time()

    input_files = inputs
    try:
        from . import package

        input_files = package.flatten_rocpd_yaml_input_file(inputs, skip_auto_merge=True)
    except Exception as e:
        print(f"Import error trying to use package, fallback to use inputs: {e}")

    output_path = kwargs.get("output_path")
    output_filename = kwargs.get("output_file")
    if not output_filename.endswith(".db"):
        output_filename += ".db"
    output = os.path.join(output_path, output_filename)

    merge_sqlite_dbs(input_files, output)

    elapsed_time = time.time() - start_time

    print(f"Merge completed successfully! Output saved to: {output}")
    print(f"Time: {elapsed_time:.2f} sec")
    return str(output)


def main(argv=None) -> int:
    """Main entry point for command line execution."""

    from . import output_config

    parser = argparse.ArgumentParser(
        description="Generate merged database from rocPD databases"
    )

    required_params = parser.add_argument_group("Required options")

    required_params.add_argument(
        "-i",
        "--input",
        required=True,
        type=output_config.check_file_exists,
        nargs="+",
        help="Path to the input ROCpd database files",
    )

    process_args = add_args(parser)

    args = parser.parse_args(argv)

    merge_args = process_args(args)

    execute(args.input, **merge_args)


if __name__ == "__main__":
    main()