import os
import re
from math import inf
from datetime import UTC, datetime
from itertools import islice

import yaml
from astropy.io import fits
from ndcube import NDCube
from prefect.variables import Variable
from prefect_sqlalchemy import SqlAlchemyConnector
from sqlalchemy import or_
from sqlalchemy.orm import Session
from yaml.loader import FullLoader

from punchbowl.auto.control.db import File
from punchbowl.data import get_base_file_name, write_ndcube_to_fits, write_ndcube_to_quicklook
from punchbowl.data.punch_io import _make_provenance_hdu

DEFAULT_SCALING = (5e-13, 5e-11)

def get_database_session(get_engine=False,
                         engine_kwargs={},
                         session_kwargs={},
                         credential_name: str = None):
    """Sets up a session to connect to the MariaDB punchpipe database

    Parameters
    ----------
    get_engine : bool
        if true returns engine as second parameter
    engine_kwargs : dict
        kwargs for the SqlAlchemy engine creation
    session_kwargs : dict
        kwargs for the SqlAlchemy session creation
    credential_name : str
        keyword for the credential block name in Prefect

    Returns
    -------
    session or (session, engine) tuple depending on ``get_engine``
    """
    if credential_name is None:
        credential_name = load_pipeline_configuration().get('credential_name', "mariadb_creds")

    credentials = SqlAlchemyConnector.load(credential_name, _sync=True)
    engine = credentials.get_engine(**engine_kwargs)
    session = Session(engine, **session_kwargs)

    if get_engine:
        return session, engine
    return session


def update_file_state(session, file_id, new_state):
    session.query(File).where(File.file_id == file_id).update({"state": new_state})
    session.commit()


def load_pipeline_configuration(path: str = None) -> dict:
    if path is None:
        path = Variable.get("punchpipe_config", "punchpipe_config.yaml")
    with open(path) as f:
        config = yaml.load(f, Loader=FullLoader)
    # TODO: add validation
    return config


def load_quicklook_scaling(level: str = None, product: str = None, obscode: str = None, path: str = None) -> (float, float):
    if path is None:
        path = Variable.get("punchpipe_config", "punchpipe_config.yaml")
    with open(path) as f:
        config = yaml.load(f, Loader=FullLoader)
    if "quicklook_scaling" in config:
        if level:
            level_data = config.get("quicklook_scaling", {}).get(level, {})
            if product and isinstance(level_data, dict):
                product_data = level_data.get(product, level_data.get("default"))
                if obscode == "4":
                    return product_data[1]
                return product_data[0]
            if obscode == "4":
                return level_data.get("default")[1]
            return level_data.get("default")[0]
        return DEFAULT_SCALING
    return DEFAULT_SCALING


def write_file(data: NDCube, corresponding_file_db_entry, pipeline_config) -> None:
    output_filename = os.path.join(
        corresponding_file_db_entry.directory(pipeline_config["root"]), corresponding_file_db_entry.filename(),
    )
    output_dir = os.path.dirname(output_filename)
    os.makedirs(output_dir, exist_ok=True)
    write_ndcube_to_fits(data,
                         output_filename,
                         write_hash=pipeline_config.get("write_sha_files", True))

    if pipeline_config.get('write_quicklooks', True) and corresponding_file_db_entry.file_type[0] not in ('S', 'T'):
        _write_quicklook(pipeline_config, corresponding_file_db_entry, data)
    return output_filename


def _write_quicklook(pipeline_config: dict, corresponding_file_db_entry: File, data: NDCube):
        ql_directory = pipeline_config.get("ql_root", pipeline_config["root"])
        ql_filename = os.path.join(corresponding_file_db_entry.directory(ql_directory),
                                   corresponding_file_db_entry.filename())
        ql_filename = ql_filename.replace(".fits", ".jp2")
        os.makedirs(os.path.dirname(ql_filename), exist_ok=True)
        write_ndcube_to_quicklook(data, ql_filename, layer="tB", write_hash=True)


def match_data_with_file_db_entry(data: NDCube, file_db_entry_list):
    # figure out which file_db_entry this corresponds to
    matching_entries = [
        file_db_entry
        for file_db_entry in file_db_entry_list
        if file_db_entry.filename() == get_base_file_name(data) + ".fits"
    ]
    if len(matching_entries) == 0:
        for file in file_db_entry_list:
            raise RuntimeError(f"There did not exist a file_db_entry for this output cube: "
                               f"result={get_base_file_name(data)}. Candidate: {file.filename()}")
    elif len(matching_entries) > 1:
        raise RuntimeError("There were many database entries matching this result. There should only be one.")
    else:
        return matching_entries[0]


def get_files_in_time_window(level: str,
                             file_type: str,
                             obs_code: str,
                             start_time: datetime,
                             end_time: datetime,
                             session: Session | None) -> list[File]:
    if session is None:
        get_database_session()

    return (session.query(File).filter(or_(File.state == "created", File.state == "progressed"))
            .filter(File.level == level)
            .filter(File.file_type == file_type)
            .filter(File.observatory == obs_code)
            .filter(File.date_obs > start_time)
            .filter(File.date_obs <= end_time).all())


def batched(iterable, n):
    # batched('ABCDEFG', 3) → ABC DEF G
    # This is basically itertools.batched, but that only exists in Python >= 3.12
    if n < 1:
        raise ValueError("n must be at least one")
    iterator = iter(iterable)
    while batch := tuple(islice(iterator, n)):
        yield batch


def group_files_by_time(files: list[File],
                        max_duration_seconds: float = inf,
                        max_per_group: int = inf) -> list[list[File]]:
    # We need to group up files by date_obs, but we need to handle small variations in date_obs. The files are coming
    # from the database already sorted, so let's just walk through the list of files and cut a group boundary every time
    # date_obs increases by more than a threshold.
    grouped_files = []
    # We'll keep track of where the current group started, and then keep stepping to find the end of this group.
    group_start = 0
    tstamp_start = files[0].date_obs.replace(tzinfo=UTC).timestamp()
    file_under_consideration = 0
    while True:
        file_under_consideration += 1
        if file_under_consideration == len(files):
            break
        this_tstamp = files[file_under_consideration].date_obs.replace(tzinfo=UTC).timestamp()
        if (abs(this_tstamp - tstamp_start) > max_duration_seconds
                or file_under_consideration - group_start >= max_per_group):
            # date_obs has jumped by more than our tolerance, so let's cut the group and then start tracking the next
            # one
            grouped_files.append(files[group_start:file_under_consideration])
            group_start = file_under_consideration
            tstamp_start = this_tstamp
    grouped_files.append(files[group_start:])
    return grouped_files


def replace_version(pattern, replacement, string):
    return re.sub(fr"_v{pattern}\.fits", f"_v{replacement}.fits", string)


def replace_file_version_in_metadata(path, old_pattern, new_version):
    with fits.open(path, mode='update', disable_image_compression=True) as hdul:
        for i, hdu in enumerate(hdul):
            if 'EXTNAME' not in hdu.header:
                continue
            if hdu.header['EXTNAME'] in ('PRIMARY DATA ARRAY', 'UNCERTAINTY ARRAY'):
                for key in hdu.header:
                    if isinstance(hdu.header[key], str):
                        hdu.header[key] = replace_version(old_pattern, new_version, hdu.header[key])
                if "FILEVRSN" in hdu.header:
                    hdu.header['FILEVRSN'] = new_version
                if "HISTORY" in hdu.header:
                    hist = hdu.header['HISTORY']
                    for i in range(len(hist)):
                        hist[i] = replace_version(old_pattern, new_version, hist[i])
            elif hdu.header['EXTNAME'] == 'FILE PROVENANCE':
                source_files = hdu.data['provenance']
                source_files = [replace_version(old_pattern, new_version, f) for f in source_files]
                hdu_provenance = _make_provenance_hdu(source_files)
                hdul[i] = hdu_provenance
