import pickle
import zipfile
import os 
import json_tricks as json

from urllib.parse import urlparse, parse_qs
from matchms import Spectrum
from matchms.exporting import save_spectra

from .mass_functions import flatten_tuplelist

class SimpleSample:
    '''
    Lightweight class of an experimental sample to facilitate workflow.
    The primary use of this class is managing data stroage and retreival at sample level.

    Depending on database_mode, sample list_mass_tracks are stored in memory, or on disk, or in MongoDB.
    Function to get mass tracks from a mzML file is in workflow.process_project and batch_EIC_from_samples_.
    Peaks and empCpds are determined in constructors.CompositeMap.
    '''
    def __init__(self, registry={}, experiment=None, database_mode='ondisk', mode='pos', is_reference=False):
        '''
        Build a lightweight sample class.

        Parameters
        ----------
        registry : dict
            sample_registry, a dictionary like {'sample_id': ii, 'input_file': file}, 
            generated by workflow.register_samples.
        experiment : ext_Experiment instance
            mostly required pointer to ext_Experiment instance, in order to get parameters.
        database_mode : str, optional, default: 'ondisk
            'ondisk' or 'memory' (run in memory, only small studies).
        mode: str, optional, default: 'pos'
            ionization mode, 'pos' or 'neg'. This should be consistent with experiment and registry if given.

        Note
        ----
        m/z calibration is performed in m/z alignment for "small studies", 
        where mass tracks are assembled to MassGrid via landmark peaks with m/z calibration per sample.
        For larger studies, m/z alignment is done via NN clustering, where m/z accuracy 
        is not checked during MassGrid construction. But it is checked during DB annotation.
        '''
        self.experiment = experiment
        self.mode = mode
        self.database_mode = database_mode 
        self.is_reference = is_reference 
        self.__registry = registry

        self.input_file = registry['input_file']
        self.name = registry['name']
        self.sample_id = registry['sample_id']
        self.data_location = registry['data_location']
        self.track_mzs = registry['track_mzs']
        self.max_scan_number = registry['max_scan_number']
        self.anchor_mz_pairs = registry['anchor_mz_pairs']
        self.rt_numbers = registry['list_scan_numbers']
        self.list_retention_time = registry['list_retention_time']
        self.compressed = self.experiment.parameters['compress']

        if self.database_mode == 'memory':
            self.list_mass_tracks = registry['sample_data']['list_mass_tracks']
        else:
            self.list_mass_tracks = []
            
        self._mz_landmarks_ = flatten_tuplelist(self.anchor_mz_pairs)
        self.rt_landmarks = []  # to populate at CMAP.calibrate_sample_RT

        # These are critical RT calibration functions, index mapping with the reference sample
        self.rt_cal_dict = None
        self.reverse_rt_cal_dict = None
        self.is_rt_aligned = is_reference      # init value False unless is_reference
        
        # placeholder
        self.mz_calibration_function = None

    @property
    def list_scan_numbers(self):
        return self.__registry['list_scan_numbers']


    @staticmethod
    def get_mass_tracks_for_sample(sample):
        return sample.get_masstracks_and_anchors()

    def get_masstracks_and_anchors(self):
        '''
        Retrieve list_mass_tracks for this sample if not alrady in memory.

        Returns
        ------- 
        A list of all mass tracks in this sample.

        Note
        ----
        Mass tracks are the bulk of data per sample, stored dependent on database_mode.
        list_mass_tracks is accessed twice in this version of asari:
        1) RT calibration and building composite map
        2) extraction of peak areas for features
        '''
        if self.list_mass_tracks:     # important, this is used to check if in memory
            return self.list_mass_tracks
        else:
            sample_data = self._get_sample_data()
            list_mass_tracks = sample_data['list_mass_tracks']
            return list_mass_tracks


    def get_rt_calibration_records(self):
        '''
        Returns a dictionary of sample_id, name, 
        rt_landmarks (list of apex scan numbers for the peaks used in RT calibration), 
        reverse_rt_cal_dict (key=reference scan number, value=sample specific scan number).
        '''
        return {
            'sample_id': self.sample_id,
            'name': self.name,
            'rt_landmarks': self.rt_landmarks,
            'reverse_rt_cal_dict': self.reverse_rt_cal_dict,
        }

    def extract_ms2(self, export_format="msp"):
        '''
        Extract MS2 data from sample data 
        '''
        try:
            if self.database_mode == 'memory':
                ms2_data = self.__registry['sample_data']['ms2_spectra']
            else:
                ms2_data = self._get_sample_data()['ms2_spectra']
            spectra = []
            for spec in ms2_data:
                mzs = []
                intensities = []    
                rtime = spec.scan_time_in_minutes()*60
                for mz, intensity in zip(spec.mz, spec.intensity):
                    mzs.append(mz)
                    intensities.append(intensity)
                try:
                    precursor_mz = spec.precursor_mz
                except:
                    precursor_mz = None
                spectra.append(Spectrum(mz=mzs, 
                                        intensities=intensities, 
                                        metadata={'scan_time': rtime,
                                                'origin': self.name,
                                                'precursor_mz': precursor_mz,
                                                }))
            if export_format[0] == ".":
                export_format = export_format[1:]
            self.experiment.parameters['ms2_export_format'] = export_format
            path = os.path.join(self.experiment.parameters['ms2_spectra_outdir'], "ms2_{}.{}".format(self.name, export_format))
            save_spectra(spectra, path, export_style="matchms")
        except Exception as _:
            print(f"Error Extracting MS2 for: {self.name}")

    def _get_sample_data(self):
        '''
        Wrapper of _retrieve_from_disk function.
        Old function kept to leave room for additional logics.

        Note:
            Potential use of database_mode, e.g.
            if self.database_mode == 'ondisk': 
                return self._retrieve_from_disk()
            elif: self.database_mode == 'firebase': 
                return self.retrieve_from_db()
        '''
        return SimpleSample.load_intermediate(self.data_location)
    
    def _retrieve_from_disk(self):
        return SimpleSample.load_intermediate(self.data_location)

    @staticmethod
    def _load_from_database(db_path, sample_id):
        """Loads sample data by querying the SQLite database."""
        raise Exception("Not Implemented")
        import sqlite3
        import pandas as pd
        if not os.path.exists(db_path):
            raise FileNotFoundError(f"Database file not found: {db_path}")

        with sqlite3.connect(db_path) as conn:
            # 1. Retrieve sample metadata
            query = "SELECT * FROM samples WHERE sample_id = ?"
            # Use pandas for easy row-to-dict conversion
            df_meta = pd.read_sql_query(query, conn, params=(sample_id,))
            if df_meta.empty:
                raise ValueError(f"No data found for sample_id {sample_id} in {db_path}")
            
            # Convert the first row to a dictionary
            sample_data = df_meta.iloc[0].to_dict()

            # 2. Deserialize JSON fields back to Python objects
            for key, value in sample_data.items():
                if isinstance(value, str) and value.startswith(('[', '{')):
                    try:
                        sample_data[key] = json.loads(value)
                    except json.JSONDecodeError:
                        pass # Keep as string if not valid JSON

            # 3. Retrieve and reconstruct mass tracks
            query_tracks = "SELECT track_id, mz, scan_number, intensity FROM mass_tracks WHERE sample_id = ?"
            df_tracks = pd.read_sql_query(query_tracks, conn, params=(sample_id,))
            
            list_mass_tracks = []
            if not df_tracks.empty:
                df_tracks = df_tracks.sort_values(['track_id', 'scan_number'])
                grouped = df_tracks.groupby(['track_id', 'mz'])['intensity'].apply(list).reset_index()
                list_mass_tracks = [
                    {'id_number': row.track_id, 'mz': row.mz, 'intensity': row.intensity}
                    for row in grouped.itertuples()
                ]
            
            sample_data['list_mass_tracks'] = list_mass_tracks
            
        return sample_data

    @staticmethod
    def _load_from_file(file_path):
        """Loads sample data from a local pickle or json file."""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Data file not found: {file_path}")

        sample_data = None
        if zipfile.is_zipfile(file_path):
            with zipfile.ZipFile(file_path, 'r') as z:
                filename_in_zip = z.namelist()[0]
                with z.open(filename_in_zip) as f:
                    if filename_in_zip.endswith('.pickle'):
                        sample_data = pickle.load(f)
                    elif filename_in_zip.endswith('.json'):
                        sample_data = json.loads(f.read().decode('utf-8'))
        else:
            if file_path.endswith('.pickle'):
                with open(file_path, 'rb') as f:
                    sample_data = pickle.load(f)
            elif file_path.endswith('.json'):
                with open(file_path, 'r') as f:
                    sample_data = json.load(f)
        
        if sample_data is None:
            raise ValueError(f"Failed to load sample data from: {file_path}")
        return sample_data

    @staticmethod
    def load_intermediate(data_location_uri):
        """
        Retrieves sample data from its storage location based on the URI scheme.

        Parameters
        ----------
        data_location_uri : str
            A URI specifying the data location, e.g.,
            'file:///path/to/data.pickle' or 
            'sqlite:///path/to/db.sqlite?sample_id=101'.
        """
        #print(f"Loading intermediate from: {data_location_uri}")
        try:
            parsed_uri = urlparse(data_location_uri)
            
            if parsed_uri.scheme == 'file':
                # Path for local files might need adjustment for Windows drive letters
                file_path = os.path.abspath(parsed_uri.netloc + parsed_uri.path)
                return SimpleSample._load_from_file(file_path)
                
            elif parsed_uri.scheme == 'sqlite':
                raise NotImplementedError
                db_path = os.path.abspath(parsed_uri.path)
                query_params = parse_qs(parsed_uri.query)
                if 'sample_id' not in query_params:
                    raise ValueError("Missing 'sample_id' in sqlite URI query")
                sample_id = int(query_params['sample_id'][0])
                return SimpleSample._load_from_database(db_path, sample_id)
                
            else:
                raise ValueError(f"Unsupported storage scheme: '{parsed_uri.scheme}'")
        except Exception as e:
            print(f"Error loading data from {data_location_uri}: {e}")
            raise