# Copyright (c) 2025 Victor Giannakouris
#
# This file is part of SwellDB and is licensed under the MIT License.
# See the LICENSE file in the project root for more information.

import json

from typing import List, Set, Dict, Union

from swelldb.engine.execution_engine import ExecutionEngine
from swelldb.llm.abstract_llm import AbstractLLM
from swelldb.table_plan import planner_prompts
from swelldb.table_plan.swelldb_schema import SwellDBSchema

from swelldb.table_plan.table.logical.logical_table import LogicalTable
from swelldb.table_plan.table.physical.dataset_table import DatasetTable
from swelldb.table_plan.table.physical.llm_table import LLMTable
from swelldb.table_plan.table.physical.physical_table import PhysicalTable
from swelldb.table_plan.table.physical.search_engine_table import SearchEngineTable


class TableGenPlanner:
    def __init__(
        self, llm: AbstractLLM, execution_engine: ExecutionEngine, serper_api_key: str
    ):
        self._llm: AbstractLLM = llm
        self._execution_engine: ExecutionEngine = execution_engine
        self._serper_api_key: str = serper_api_key

    def create_plan_from_operators(
        self,
        logical_table: LogicalTable,
        base_columns: List[str],
        tables: Dict[str, str],
        operators: List[type],
    ):
        # The initial column set, defined by the user
        initial_schema: SwellDBSchema = logical_table.get_schema()

        # The remaining column set — Keeps track of the columns that are not yet generated
        remaining_column_set: Set[str] = set(initial_schema.get_attribute_names())

        root = None
        while operators and remaining_column_set:
            curr_logical_table_schema = SwellDBSchema(
                attributes=[
                    logical_table.get_schema().get_attribute(col)
                    for col in remaining_column_set
                ]
            )

            curr_logical_table = LogicalTable(
                name="tbl",
                prompt=logical_table.get_prompt(),
                schema=curr_logical_table_schema,
            )

            operator_cls = operators.pop(0)
            prompt: str = operator_cls.get_columns_prompt(curr_logical_table, tables)

            if not prompt:
                operator_columns = list(remaining_column_set)
            else:
                operator_columns = [
                    c.strip() for c in self._llm.call(prompt).split(",")
                ]

            for base_column in base_columns:
                if base_column not in operator_columns:
                    operator_columns.append(base_column)

            new_logical_table_schema = SwellDBSchema(
                attributes=[
                    logical_table.get_schema().get_attribute(col)
                    for col in operator_columns
                ]
            )

            new_logical_table = LogicalTable(
                name="tbl",
                prompt=logical_table.get_prompt(),
                schema=new_logical_table_schema,
            )

            # Replace the root table with the new operator — Add the previous root as its child.
            root = operator_cls(
                self._execution_engine, new_logical_table, root, base_columns, self._llm
            )

            remaining_column_set.difference_update(operator_columns)

        return root

    def create_plan(
        self,
        logical_table: LogicalTable,
        base_columns: List[str],
        tables: Dict[str, str] = dict(),
    ) -> PhysicalTable:

        # Root table operator
        root_table: PhysicalTable = None

        # The initial column set, defined by the user
        initial_schema: SwellDBSchema = logical_table.get_schema()

        # The remaining column set — Keeps track of the columns that are not yet generated
        remaining_column_set: Set[str] = set(initial_schema.get_attribute_names())

        # Get the columns that we can generate with the provided datasets
        local_ds_columns_response: str = self._llm.call(
            planner_prompts.get_local_tables_prompt(
                logical_table, table_schema_dict=tables
            )
        )

        # This response is a JSON string that contains the columns and the SQL query
        local_ds_columns_meta: Dict = json.loads(local_ds_columns_response)
        local_ds_columns = local_ds_columns_meta["columns"]
        local_ds_sql_query = local_ds_columns_meta["query"]

        # If any of the columns can be generated by the LLM
        if local_ds_columns_meta:
            local_ds_table_schema = SwellDBSchema(
                attributes=[
                    logical_table.get_schema().get_attribute(col)
                    for col in local_ds_columns
                ]
            )

            local_ds_logical_table = LogicalTable(
                name="search_table",
                prompt=logical_table.get_prompt(),
                schema=local_ds_table_schema,
            )

            # Set the root table
            root_table = DatasetTable(
                logical_table=local_ds_logical_table,
                llm=self._llm,
                base_columns=base_columns,
                child_table=None,
                query=local_ds_sql_query,
                execution_engine=self._execution_engine,
            )

            # Remove the columns that can be generated with the provided datasets from the remaining column set
            remaining_column_set = remaining_column_set.difference(local_ds_columns)

        # Check which columns can be generated by the LLM
        if remaining_column_set:
            tmp_logical_table = LogicalTable(
                name="llm_table",
                prompt=logical_table.get_prompt(),
                schema=SwellDBSchema(
                    attributes=[
                        logical_table.get_schema().get_attribute(col)
                        for col in remaining_column_set
                    ]
                ),
            )

            # Columns that the LLM can generate
            llm_columns: str = self._llm.call(
                planner_prompts.get_llm_columns_prompt(tmp_logical_table)
            )

            print(f"LLM columns: {llm_columns}")

            if llm_columns:
                llm_columns_list: Set[str] = set(
                    [x.strip() for x in llm_columns.split(",")]
                )

                remaining_column_set.difference_update(llm_columns_list)

                # add all base columns to llm_columns
                for base_column in base_columns:
                    if base_column not in llm_columns_list:
                        llm_columns_list.add(base_column)

                llm_table_schema: SwellDBSchema = SwellDBSchema(
                    attributes=[
                        logical_table.get_schema().get_attribute(col)
                        for col in llm_columns_list
                    ]
                )

                llm_logical_table = LogicalTable(
                    name="llm_table",
                    prompt=logical_table.get_prompt(),
                    schema=llm_table_schema,
                )

                root_table = LLMTable(
                    llm=self._llm,
                    execution_engine=self._execution_engine,
                    logical_table=llm_logical_table,
                    child_table=root_table,
                    base_column=base_column,
                )

        if remaining_column_set:
            remaining_column_set.add(base_column)

            search_table_schema = SwellDBSchema(
                attributes=[
                    logical_table.get_schema().get_attribute(col)
                    for col in remaining_column_set
                ]
            )

            search_logical_table = LogicalTable(
                name="search_table",
                prompt=logical_table.get_prompt(),
                schema=search_table_schema,
            )

            root_table = SearchEngineTable(
                llm=self._llm,
                logical_table=search_logical_table,
                child_table=root_table,
                base_column=base_column,
                execution_engine=self._execution_engine,
                serper_api_key=self._serper_api_key,
            )

        return root_table
