# -*- coding: utf-8 -*-
import os
import json
import shutil
import tempfile
import psycopg2
from psycopg2 import sql
from pydantic import BaseModel
from psycopg2.extras import RealDictCursor
from dagstermill.factory import execute_notebook
from typing import Callable, Iterable, Optional, Mapping, Union, Any, List, Set, Tuple
from dagster import define_asset_job, Shape, Field, AssetIn, Output, AssetKey, AssetsDefinition, JobDefinition

# Importing from private dagster classes (DagsterMill)
import dagster._check as check
from dagster._core.execution.context.compute import OpExecutionContext

# Oxaigen SDK
from oxaigen.src.settings.settings import BaseSettingsInterface

# Internal Imports
from ..constant import DEFAULT_REPORT_FILE_TYPE
from ..settings.settings import Settings
from ..notebook.notebook_document_generator import render_dagstermill_notebook_to_output_file_type
from ..notebook.s3_storage import S3ClientStorage
from .notebook_asset import define_notebook_as_asset

# CONSTANTS
DEFAULT_REPORT_ASSET_PREFIX: List[str] = ["Oxaigen", "Reports"]
DEFAULT_REPORT_ASSET_GROUP_NAME = "Reports"
DEFAULT_REPORT_IO_MANAGER = "reports_io_manager"
DEFAULT_S3_STORAGE_PREFIX = "reports"
WORKSPACE_ID = "workspaceId"
REPORT_FILE_TYPE = "reportFileType"
FILES_DIR = "files"
CLEAN_FILES_DIR = "clean"


class S3StorageMetadata(BaseModel):
    """
    Pydantic model for S3 storage metadata.
    """
    PermanentNotebookLink: str
    PermanentOutputFileLink: str


def _execute_notebook_step(
        context: OpExecutionContext,
        name: str,
        notebook_path: str,
        save_notebook_on_failure: bool,
        settings: BaseSettingsInterface,
        **inputs
) -> tuple[str, str]:
    """
    Executes a Jupyter notebook as part of a Dagster step, saves both the executed
    and a cleaned version of the notebook to disk, and returns their file paths.

    The executed notebook is saved in the reports directory. A cleaned version is
    also created by removing the first and last cells, typically used to strip
    metadata or boilerplate from automated runs.

    Args:
        context (OpExecutionContext): The Dagster execution context for the current step.
        name (str): A name identifier for the notebook execution run.
        notebook_path (str): Path to the input notebook to be executed.
        save_notebook_on_failure (bool): Whether to save the notebook even if execution fails.
        settings (BaseSettingsInterface): Settings instance containing configuration
            such as the output directory path. Defaults to the global `settings` object.
        **inputs: Arbitrary keyword arguments passed as parameters to the executed notebook.

    Returns:
        tuple[str, str]: A tuple containing:
            - The path to the executed notebook file.
            - The path to the cleaned notebook file (with first and last cells removed).
    """
    with tempfile.TemporaryDirectory() as output_notebook_dir:
        executed_notebook_path = execute_notebook(
            context.get_step_execution_context(),
            name=name,
            inputs=inputs,
            save_notebook_on_failure=save_notebook_on_failure,
            notebook_path=notebook_path,
            output_notebook_dir=output_notebook_dir,
        )

        # Paths
        reports_file_dir = os.path.join(settings.REPORTS_DIR, FILES_DIR)
        clean_subdir = os.path.join(reports_file_dir, CLEAN_FILES_DIR)

        os.makedirs(reports_file_dir, exist_ok=True)
        os.makedirs(clean_subdir, exist_ok=True)

        notebook_base_name = os.path.basename(executed_notebook_path)
        copied_notebook_path = os.path.join(reports_file_dir, notebook_base_name)
        cleaned_notebook_path = os.path.join(clean_subdir, notebook_base_name)

        # Copy the original executed notebook
        shutil.copy(executed_notebook_path, copied_notebook_path)

        # Create a cleaned version (remove first and last cell including output)
        with open(executed_notebook_path, "r", encoding="utf-8") as f:
            nb = json.load(f)

        if "cells" in nb and len(nb["cells"]) >= 3:
            nb["cells"] = nb["cells"][1:-1]

        with open(cleaned_notebook_path, "w", encoding="utf-8") as f:
            json.dump(nb, f, ensure_ascii=False, indent=2)

        return copied_notebook_path, cleaned_notebook_path


def _generate_output_file_step(
        executed_notebook_path: str,
        settings: BaseSettingsInterface,
        file_type: Optional[str] = DEFAULT_REPORT_FILE_TYPE
) -> str:
    """
    Converts an executed Jupyter notebook into a PDF report.

    This step cleans the notebook using a predefined cleaning function,
    then renders the cleaned version to PDF using a notebook-to-PDF renderer.
    The resulting PDF is saved to the reports' directory.

    Args:
        executed_notebook_path (str): Path to the executed Jupyter notebook file.
        settings (BaseSettingsInterface): Settings instance containing configuration
            such as the reports directory path. Defaults to the global `settings` object.
        file_type (str): Output file type (PDF/DocX/etc.)

    Returns:
        str: The file path to the generated PDF report.

    Raises:
        FileNotFoundError: If the notebook cleaning or PDF rendering fails at any step.
    """
    reports_dir = settings.REPORTS_DIR
    reports_file_dir = os.path.join(reports_dir, FILES_DIR)
    clean_reports_file_dir = os.path.join(reports_file_dir, CLEAN_FILES_DIR)

    notebook_base_name = os.path.basename(executed_notebook_path)

    file_generated: bool = render_dagstermill_notebook_to_output_file_type(
        working_directory=reports_dir,
        executed_notebook_file_name=notebook_base_name,
        file_type=file_type.lower()
    )
    if file_generated is False:
        raise FileNotFoundError(f"Could not generate {file_type} file from clean report notebook")

    output_file_path = os.path.join(clean_reports_file_dir, notebook_base_name.replace(".ipynb", f".{str(file_type).lower() }"))

    return output_file_path


def _store_files_to_s3_step(
        clean_notebook_path: str,
        output_file_path: str,
        run_id: str,
        name: str,
        settings: BaseSettingsInterface,
        asset_key_prefix: Optional[List[str]] = DEFAULT_REPORT_ASSET_PREFIX,
        s3_prefix: Optional[Union[str, List[str]]] = DEFAULT_S3_STORAGE_PREFIX,
        output_file_type: Optional[str] = DEFAULT_REPORT_FILE_TYPE,
) -> S3StorageMetadata:
    """
    Uploads the cleaned notebook and corresponding PDF report to S3 and returns metadata.

    This step stores both the cleaned `.ipynb` file and the rendered `.pdf` file
    to structured paths in an S3 bucket, organized by `run_id` and `name`. The returned
    metadata includes permanent S3 links to both files.

    Args:
        clean_notebook_path (str): Local path to the cleaned Jupyter notebook file.
        output_file_path (str): Local path to the generated report output file.
        run_id (str): Unique identifier for the notebook execution run.
        name (str): Name of the notebook, used in constructing the S3 filenames.
        settings (BaseSettingsInterface): Settings instance containing S3 configuration.
            Defaults to the global `settings` object.
        asset_key_prefix (List[str], Optional): the S3 Dagster Bucket directory path for the notebook asset
        s3_prefix (Union[str, List[str]]], Optional): the S3 Client Bucket directory path where the notebook and pdf will be stored.

    Returns:
        S3StorageMetadata: An object containing permanent S3 links for the uploaded notebook and PDF.
    """
    bucket = settings.DAGSTER_BUCKET if settings.DAGSTER_BUCKET is not None else settings.S3_BUCKET
    s3_storage = S3ClientStorage(settings=settings, bucket=bucket)

    # Normalize s3_prefix to a list
    if isinstance(s3_prefix, str):
        s3_prefix_parts = [s3_prefix]
    elif isinstance(s3_prefix, list):
        s3_prefix_parts = s3_prefix
    else:
        s3_prefix_parts = []

    # Construct full S3 path prefix
    s3_file_path = os.path.join(*(s3_prefix_parts + asset_key_prefix))
    notebook_s3_path = os.path.join(s3_file_path, run_id, f"{name}.ipynb")

    # Save the notebook
    s3_storage.save_file(file_path=clean_notebook_path, s3_path=notebook_s3_path)

    # Save the Output File
    output_s3_file_path = os.path.join(s3_file_path, run_id, f"{name}.{output_file_type.lower()}")
    s3_storage.save_file(file_path=output_file_path, s3_path=output_s3_file_path)

    return S3StorageMetadata(
        PermanentNotebookLink=f"S3://{settings.DAGSTER_BUCKET}/{notebook_s3_path}",
        PermanentOutputFileLink=f"S3://{settings.DAGSTER_BUCKET}/{output_s3_file_path}"
    )


def check_reports_table_exists(
        settings: BaseSettingsInterface,
        schema_name: str = "reports",
        table_name: str = "Report"
) -> bool:
    """
    Checks if the Reports table exists in the specified PostgreSQL schema.

    Args:
        settings (BaseSettingsInterface): Settings instance providing access to the
            database connection URI.
        schema_name (str, optional): Name of the schema to check. Defaults to "reports".
        table_name (str, optional): Name of the table to check. Defaults to "Report".

    Returns:
        bool: True if the table exists, False otherwise.

    Raises:
        Exception: If any database error occurs during the check.
    """
    conn = None
    try:
        conn = psycopg2.connect(settings.PSYCOPG_CLIENT_DATABASE_URI)
        cur = conn.cursor()

        check_query = sql.SQL("""
                              SELECT EXISTS (SELECT
                                             FROM information_schema.tables
                                             WHERE table_schema = %s
                                               AND table_name = %s)
                              """)

        cur.execute(check_query, (schema_name, table_name))
        exists = cur.fetchone()[0]

        cur.close()
        conn.close()

        return exists

    except Exception as e:
        print(f"Error checking if reports table exists: {e}")
        if conn:
            conn.close()
        raise


def create_reports_table(
        settings: BaseSettingsInterface,
        schema_name: str = "reports",
        table_name: str = "Report"
) -> None:
    """
    Creates the Reports table in the specified PostgreSQL schema if it doesn't exist.

    Args:
        settings (BaseSettingsInterface): Settings instance providing access to the
            database connection URI.
        schema_name (str, optional): Name of the schema to create the table in.
            Defaults to "reports".
        table_name (str, optional): Name of the table to create. Defaults to "Report".

    Raises:
        Exception: If any database error occurs during table creation.
    """
    conn = None
    try:
        conn = psycopg2.connect(settings.PSYCOPG_CLIENT_DATABASE_URI)
        cur = conn.cursor()

        # Create schema if it doesn't exist
        create_schema_query = sql.SQL("CREATE SCHEMA IF NOT EXISTS {}").format(
            sql.Identifier(schema_name)
        )
        cur.execute(create_schema_query)

        # Enable uuid extension if not already enabled
        cur.execute("CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\"")

        # Create table with proper structure matching your schema
        create_table_query = f'''
            CREATE TABLE IF NOT EXISTS "{schema_name}"."{table_name}" (
                "itemId" BIGSERIAL PRIMARY KEY,
                "keyPrefix" TEXT[] NOT NULL,
                "reportName" VARCHAR(255) NOT NULL,
                "reportDescription" TEXT,
                "notebookLink" VARCHAR(255) NOT NULL,
                "outputFileLink" VARCHAR(255) NOT NULL,
                "runId" UUID NOT NULL,
                "configSchema" JSONB NOT NULL,
                "createdTimestamp" TIMESTAMP WITH TIME ZONE DEFAULT NOW() NOT NULL,
                "status" SMALLINT DEFAULT 0 NOT NULL,
                "metadataUuid" UUID DEFAULT uuid_generate_v4() NOT NULL
            )
        '''
        cur.execute(create_table_query)

        conn.commit()
        cur.close()
        conn.close()

        print(f"✓ Successfully created table {schema_name}.{table_name}")

    except Exception as e:
        print(f"Error creating reports table: {e}")
        if conn:
            conn.rollback()
            conn.close()
        raise


def _insert_report_record(
        asset_key_prefix: List[str],
        name: str,
        notebook_link: str,
        output_file_link: str,
        run_id: str,
        config_dict: dict,
        settings: BaseSettingsInterface,
) -> int:
    """
    Inserts a new report entry into the reports database table and returns the generated record ID.

    The inserted record includes metadata such as the report name, S3 links to the notebook and
    report output file (PDF/DocX/etc.), run identifier, configuration schema, and a default status.

    Args:
        asset_key_prefix (List[str]): A list of string components representing the logical key prefix
            for the report (e.g. hierarchical folder-like structure).
        name (str): The name of the report or notebook.
        notebook_link (str): S3 URI pointing to the stored notebook file.
        output_file_link (str): S3 URI pointing to the stored report output file (PDF/DocX/etc.).
        run_id (str): Unique identifier for the report execution run.
        config_dict (dict): Dictionary containing the serialized configuration schema used to
            generate the report.
        settings (BaseSettingsInterface): Settings instance providing access to the
            database connection URI. Defaults to the global `settings` object.

    Returns:
        int: The primary key (`itemId`) of the inserted report record.

    Raises:
        Exception: If any database error occurs during insertion, the error is printed and re-raised.
    """
    conn = None
    try:
        # Check if table exists, create if it doesn't
        if not check_reports_table_exists(settings):
            create_reports_table(settings)

        conn = psycopg2.connect(settings.PSYCOPG_CLIENT_DATABASE_URI)
        cur = conn.cursor(cursor_factory=RealDictCursor)

        status: int = 0

        insert_query = '''
            INSERT INTO "reports"."Report" ("keyPrefix", "reportName", "notebookLink", "outputFileLink",
                                           "runId", "configSchema", "status")
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            RETURNING "itemId"
        '''
        cur.execute(insert_query,
                    (asset_key_prefix, name, notebook_link, output_file_link, run_id, json.dumps(config_dict), status))

        # Fetch the generated ID
        result = cur.fetchone()
        generated_id = result['itemId']

        conn.commit()
        cur.close()
        conn.close()

        return generated_id

    except Exception as e:
        print(f"Error inserting report record: {e}")
        if conn:
            conn.rollback()
        raise

def _make_oxaigen_report_asset_compute_function(
        asset_key_prefix: List[str],
        name: str,
        notebook_path: str,
        settings: BaseSettingsInterface,
        save_notebook_on_failure: Optional[bool] = False,
        s3_prefix: Optional[Union[str, List[str]]] = DEFAULT_S3_STORAGE_PREFIX,
        config_schema: Optional[Union[Any, Mapping[str, Any]]] = None,
        tags: Optional[Mapping[str, Any]] = None,
) -> Callable:
    """
    Creates a Dagster compute function that runs a notebook-based report pipeline.

    The returned function performs the following steps when executed within a Dagster
    pipeline:
        1. Executes a specified Jupyter notebook with given inputs.
        2. Generates a cleaned version of the notebook and converts it to a Report Document (PDF/DocX/etc.).
        3. Uploads the notebook and Output Report File to S3, returning permanent links.
        4. Inserts a record into the reports database table with metadata.
        5. Yields the executed notebook as a Dagster output, including S3 metadata.
        6. Cleans up all temporary files created during execution.

    Args:
        asset_key_prefix (List[str]): Logical key prefix used for organizing report records and S3 paths.
        name (str): Unique name of the report, used for naming files and DB records.
        notebook_path (str): Path to the notebook to be executed.
        settings (BaseSettingsInterface): Settings object with config for S3, DB, and directories.
            Defaults to the global `settings` object.
        save_notebook_on_failure (bool, optional): If True, save the executed notebook even on failure, defaults to False
        s3_prefix (Union[str,List[str]], Optional): the S3 Client Bucket directory path where the notebook and
            generated report file will be stored.
        config_schema (Optional[Union[Any, Mapping[str, Any]]], optional): Optional config schema for
            validating run configuration. Not directly used in this function, but can be useful
            in surrounding pipeline definitions.
        tags (Optional[Dict[str, Any]]): A dictionary of tags for the asset itself (not the op).
            These are visible in the Dagster UI. The tag `oxaigen_asset_type: report` will always
            be added and cannot be overridden.
    Returns:
        Callable: A Dagster compute function that can be used as an asset or op, yielding
            the notebook as a binary output with associated metadata.
    """

    def _t_fn(context: OpExecutionContext, **inputs) -> Iterable:
        check.param_invariant(
            isinstance(context.run_config, dict),
            "context",
            "StepExecutionContext must have valid run_config",
        )

        # Get execution context
        run_id = context.run_id
        config_dict = context.op_config

        selected_file_type = config_dict.get(REPORT_FILE_TYPE, DEFAULT_REPORT_FILE_TYPE)

        # Track paths to delete at the end
        temp_paths = []

        try:
            # Step 1: Execute the notebook
            executed_notebook_path, cleaned_notebook_path = _execute_notebook_step(
                context=context,
                name=name,
                notebook_path=notebook_path,
                save_notebook_on_failure=save_notebook_on_failure,
                settings=settings,
                **inputs
            )
            temp_paths.extend([executed_notebook_path, cleaned_notebook_path])

            # Step 2: Generate the output file (PDF/DocX/etc.)
            output_file_path = _generate_output_file_step(
                executed_notebook_path=cleaned_notebook_path,
                settings=settings,
                file_type=selected_file_type
            )
            temp_paths.append(output_file_path)

            # Step 3: Store files to S3
            s3_storage_metadata: S3StorageMetadata = _store_files_to_s3_step(
                clean_notebook_path=cleaned_notebook_path,
                output_file_path=output_file_path,
                run_id=run_id,
                name=name,
                settings=settings,
                asset_key_prefix=asset_key_prefix,
                s3_prefix=s3_prefix,
                output_file_type=selected_file_type
            )

            # Step 4: Insert a new record into the reports table
            _insert_report_record(
                asset_key_prefix=asset_key_prefix,
                name=name,
                settings=settings,
                notebook_link=s3_storage_metadata.PermanentNotebookLink,
                output_file_link=s3_storage_metadata.PermanentOutputFileLink,
                run_id=run_id,
                config_dict=config_dict
            )

            # Step 5: Yield notebook as Dagster output
            with open(executed_notebook_path, "rb") as fd:
                print(f"EXECUTED NOTEBOOK PATH: {executed_notebook_path}")
                yield Output(fd.read(), metadata={
                    "ReportUUID": run_id,
                    WORKSPACE_ID: settings.WORKSPACE_ID,
                    **s3_storage_metadata.model_dump(),
                })

        finally:
            # Step 6: Cleanup temporary files
            for path in temp_paths:
                try:
                    os.remove(path)
                except Exception as e:
                    print(f"⚠️ Could not delete temp file: {path} — {e}")

    return _t_fn


def define_oxaigen_report_as_asset(
        name: str,
        file: str,
        description: str,
        settings: Optional[BaseSettingsInterface] = None,
        group: Optional[str] = DEFAULT_REPORT_ASSET_GROUP_NAME,
        asset_key_prefix: Optional[List[str]] = DEFAULT_REPORT_ASSET_PREFIX,
        s3_prefix: Optional[Union[str, List[str]]] = DEFAULT_S3_STORAGE_PREFIX,
        ins: Optional[Mapping[str, AssetIn]] = None,
        config_schema: Optional[Union[Any, Mapping[str, Any]]] = None,
        tags: Optional[Mapping[str, Any]] = None,
        required_resource_keys: Optional[Set[str]] = None,
        io_manager_key: Optional[str] = DEFAULT_REPORT_IO_MANAGER
) -> Tuple[AssetsDefinition, JobDefinition]:
    """
    Defines a Dagster asset for executing a Jupyter notebook report within the Oxaigen framework.

    This function wraps the report notebook execution pipeline into a reusable Dagster asset,
    which includes notebook execution, PDF generation, file storage on S3, and database
    record insertion. The returned asset integrates with Dagster’s IO manager and resource system.

    Args:
        name (str): Unique name for the report asset.
        file (str): Path to the Jupyter notebook file to execute.
        description (str): Description of the report asset for documentation purposes.
        settings (Optional[BaseSettingsInterface], optional): Configuration settings object.
            Defaults to the global `settings` instance.
        group (Optional[str], optional): Asset group name for organizing assets in Dagster.
            Defaults to `REPORT_ASSET_GROUP_NAME`.
        asset_key_prefix (Optional[List[str]], optional): Key prefix used to namespace the asset keys.
            Defaults to `REPORT_ASSET_PREFIX`.
        s3_prefix (Union[str,List[str]], Optional): the S3 Client Bucket directory path where the notebook and pdf will
            be stored.
        ins (Optional[Mapping[str, AssetIn]], optional): Input dependencies for the asset.
            Defaults to None.
        config_schema (Optional[Union[Any, Mapping[str, Any]]], optional): Configuration schema
            for validating run-time inputs. Defaults to None.
        tags (Optional[Dict[str, Any]]): A dictionary of tags for the asset itself (not the op).
            These are visible in the Dagster UI. The tag `oxaigen_asset_type: report` will always
            be added and cannot be overridden.
        required_resource_keys (Optional[Set[str]], optional): Set of resource keys required by the asset.
            Defaults to None.
        io_manager_key (Optional[str], optional): Key of the IO manager to use for this asset.
            Defaults to `DEFAULT_REPORT_IO_MANAGER`.

    Returns:
        Tuple[AssetDefinition, AssetJob]: A tuple containing:
            - The Dagster asset definition encapsulating the notebook report pipeline.
            - A Dagster job that runs the asset, named based on its full asset key joined by `__` and suffixed with `__job`.

    Example usage with `reports_io_manager`:

        ```python
        from dagster import Definitions
        from your_project.assets.report_assets import define_oxaigen_report_as_asset
        from oxaigen.src.io_manager.reports_io_manager import S3ReportsIOManager
        from dagster_aws.s3 import S3Resource

        s3_resource = S3Resource(...)
        reports_io_manager = S3ReportsIOManager(
            s3_resource=s3_resource,
            s3_bucket="your-s3-bucket-name",
            s3_prefix="reports",
            asset_key_prefix=["Oxaigen", "Reports"]
        )

        report_asset = define_oxaigen_report_as_asset(
            name="monthly_sales_report",
            file="reports/monthly_sales_report.ipynb",
            description="Generates monthly sales report notebook and PDF, uploads to S3 and logs metadata.",
            config_schema={"start_date": str, "end_date": str},
            io_manager_key="reports_io_manager"
        )

        defs = Definitions(
            assets=[report_asset],
            resources={
                "reports_io_manager": reports_io_manager,
                "s3_resource": s3_resource,
                # other resources as needed
            },
        )
        ```
    """
    if settings is None:
        settings = Settings()

    # Handle config_schema to ensure workspaceId field is present
    if config_schema is None:
        config_schema = Shape({
            REPORT_FILE_TYPE: Field(str, is_required=False, description="The document type for the generated report")
        })
    elif isinstance(config_schema, Shape):
        existing_fields = config_schema.fields.copy()
        if REPORT_FILE_TYPE not in config_schema.fields:
            existing_fields[REPORT_FILE_TYPE] = Field(str, is_required=False, description="The document type for the generated report")
        config_schema = Shape(existing_fields)
    elif not isinstance(config_schema, Shape):
        raise ValueError("Invalid config_schema, use Shape() from dagster")

    # Define report compute function (callable)
    oxaigen_report_asset_compute_function = _make_oxaigen_report_asset_compute_function(
        asset_key_prefix=asset_key_prefix,
        name=name,
        notebook_path=file,
        save_notebook_on_failure=True,
        s3_prefix=s3_prefix,
        settings=settings,
        config_schema=config_schema,
        tags=tags
    )

    # return dagster @asset wrapper asset variable
    notebook_asset = define_notebook_as_asset(
        asset_compute_function=oxaigen_report_asset_compute_function,
        key_prefix=asset_key_prefix,
        name=name,
        description=description,
        notebook_path=file,
        group_name=group,
        io_manager_key=io_manager_key,
        required_resource_keys=required_resource_keys,
        ins=ins,
        save_notebook_on_failure=True,
        config_schema=config_schema,
        tags=tags
    )

    # Construct asset key for job naming
    asset_key = AssetKey(path=(asset_key_prefix or []) + [name])
    notebook_asset_job_name = "__".join(asset_key.path) + "__job"

    # Create an asset job
    notebook_asset_job = define_asset_job(
        name=notebook_asset_job_name,
        selection=[notebook_asset.key]
    )

    return notebook_asset, notebook_asset_job
