from io import BytesIO, StringIO
import os
import csv
import json
import codecs
import traceback
import tempfile
from urllib.parse import urlparse

       
import magic
import requests
import pandas as pd
from charset_normalizer import from_bytes

from mindsdb_sql import parse_sql
from mindsdb_sql.parser.ast.base import ASTNode
from mindsdb_sql.parser.ast import DropTables, Select

from mindsdb.api.mysql.mysql_proxy.utilities.sql import query_df
from mindsdb.integrations.libs.base import DatabaseHandler
from mindsdb.integrations.libs.response import (
    HandlerStatusResponse as StatusResponse,
    HandlerResponse as Response,
    RESPONSE_TYPE
)


def clean_row(row):
    n_row = []
    for cell in row:
        if str(cell) in ['', ' ', '  ', 'NaN', 'nan', 'NA']:
            n_row.append(None)
        else:
            n_row.append(cell)

    return n_row


class FileHandler(DatabaseHandler):
    """
    Handler for files
    """
    name = 'files'

    def __init__(self, name=None, file_storage=None, connection_data={}, file_controller=None, **kwargs):
        super().__init__(name)
        self.parser = parse_sql
        self.fs_store = file_storage
        self.custom_parser = connection_data.get('custom_parser', None)
        self.clean_rows = connection_data.get('clean_rows', True)
        self.file_controller = file_controller

    def connect(self, **kwargs):
        return

    def disconnect(self, **kwargs):
        return

    def check_connection(self) -> StatusResponse:
        return StatusResponse(True)

    def query(self, query: ASTNode) -> Response:
        if type(query) == DropTables:
            for table_identifier in query.tables:
                if len(table_identifier.parts) == 2 and table_identifier.parts[0] != self.name:
                    return Response(
                        RESPONSE_TYPE.ERROR,
                        error_message=f"Can't delete table from database '{table_identifier.parts[0]}'"
                    )
                table_name = table_identifier.parts[-1]
                try:
                    self.file_controller.delete_file(table_name)
                except Exception as e:
                    return Response(
                        RESPONSE_TYPE.ERROR,
                        error_message=f"Can't delete table '{table_name}': {e}"
                    )
            return Response(RESPONSE_TYPE.OK)
        elif type(query) == Select:
            table_name = query.from_table.parts[-1]
            file_path = self.file_controller.get_file_path(table_name)
            df, _columns = self._handle_source(file_path, self.clean_rows, self.custom_parser)
            result_df = query_df(df, query)
            return Response(
                RESPONSE_TYPE.TABLE,
                data_frame=result_df
            )
        else:
            return Response(
                RESPONSE_TYPE.ERROR,
                error_message="Only 'select' and 'drop' queries allowed for files"
            )

    def native_query(self, query: str) -> Response:
        ast = self.parser(query, dialect='mindsdb')
        return self.query(ast)

    @staticmethod
    def _handle_source(file_path, clean_rows=True, custom_parser=None):
        # get file data io, format and dialect
        data, fmt, dialect = FileHandler._get_data_io(file_path)
        data.seek(0)  # make sure we are at 0 in file pointer

        if custom_parser:
            header, file_data = custom_parser(data, fmt)

        elif fmt == 'parquet':
            df = pd.read_parquet(data)
            header = df.columns.values.tolist()
            file_data = df.values.tolist()

        elif fmt == 'csv':
            df = pd.read_csv(data, sep=dialect.delimiter)
            header = df.columns.values.tolist()
            file_data = df.values.tolist()

        elif fmt in ['xlsx', 'xls']:
            data.seek(0)
            df = pd.read_excel(data)
            header = df.columns.values.tolist()
            file_data = df.values.tolist()

        elif fmt == 'json':
            data.seek(0)
            json_doc = json.loads(data.read())
            df = pd.json_normalize(json_doc, max_level=0)
            header = df.columns.values.tolist()
            file_data = df.values.tolist()

        else:
            raise ValueError('Could not load file into any format, supported formats are csv, json, xls, xlsx')

        if clean_rows:
            file_list_data = [clean_row(row) for row in file_data]
        else:
            file_list_data = file_data

        header = [x.strip() for x in header]
        col_map = dict((col, col) for col in header)
        return pd.DataFrame(file_list_data, columns=header), col_map

    @staticmethod
    def _get_data_io(file_path):
        """
        @TODO: Use python-magic to simplify the function and detect the file types as the xlsx example
        This gets a file either url or local file and defines what the format is as well as dialect
        :param file: file path or url
        :return: data_io, format, dialect
        """

        ############
        # get file as io object
        ############

        # file_path = self._get_file_path()

        data = BytesIO()

        try:
            with open(file_path, 'rb') as fp:
                data = BytesIO(fp.read())
        except Exception as e:
            error = 'Could not load file, possible exception : {exception}'.format(exception=e)
            print(error)
            raise ValueError(error)

        dialect = None

        ############
        # check for file type
        ############

        # Check first and last 4 bytes equal to PAR1.
        # Refer: https://parquet.apache.org/docs/file-format/
        parquet_sig = b'PAR1'
        data.seek(0, 0)
        start_meta = data.read(4)
        data.seek(-4, 2)
        end_meta = data.read()
        data.seek(0)
        if start_meta == parquet_sig and end_meta == parquet_sig:
            return data, 'parquet', None
        else:
            print("It's not parquet file. Checking for other formats")

        file_type = magic.from_file(file_path, mime=True)
        if file_type in ['application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', 'application/vnd.ms-excel']:
            #for now we send xlsx for all excel types since we use pd.read_excel
            return data, 'xlsx', dialect
       
        # if not excel it can be a json file or a CSV, convert from binary to stringio

        byte_str = data.read()
        # Move it to StringIO
        try:
            # Handle Microsoft's BOM "special" UTF-8 encoding
            if byte_str.startswith(codecs.BOM_UTF8):
                data = StringIO(byte_str.decode('utf-8-sig'))
            else:
                file_encoding_meta = from_bytes(
                    byte_str[:32 * 1024],
                    steps=32,           # Number of steps/block to extract from my_byte_str
                    chunk_size=1024,    # Set block size of each extraction)
                    explain=False
                )
                best_meta = file_encoding_meta.best()
                errors = 'strict'
                if best_meta is not None:
                    encoding = file_encoding_meta.best().encoding
                else:
                    encoding = 'utf-8'
                    errors = 'replace'
                data = StringIO(byte_str.decode(encoding, errors))
        except Exception:
            print(traceback.format_exc())
            print('Could not load into string')

        # see if its JSON
        buffer = data.read(100)
        data.seek(0)
        text = buffer.strip()
        # analyze first n characters
        if len(text) > 0:
            text = text.strip()
            # it it looks like a json, then try to parse it
            if text.startswith('{') or text.startswith('['):
                try:
                    json.loads(data.read())
                    data.seek(0)
                    return data, 'json', dialect
                except Exception:
                    data.seek(0)
                    return data, None, dialect

        # lets try to figure out if its a csv
        try:
            dialect = FileHandler._get_csv_dialect(data)
            if dialect:
                return data, 'csv', dialect
            return data, None, dialect
        except Exception:
            data.seek(0)
            print('Could not detect format for this file')
            print(traceback.format_exc())
            # No file type identified
            return data, None, dialect

    @staticmethod
    def _get_file_path(path) -> str:
        try:
            is_url = urlparse(path).scheme in ('http', 'https')
        except Exception:
            is_url = False
        if is_url:
            path = FileHandler._fetch_url(path)
        return path

    @staticmethod
    def _get_csv_dialect(buffer) -> csv.Dialect:
        sample = buffer.readline()  # trying to get dialect from header
        buffer.seek(0)
        try:
            if isinstance(sample, bytes):
                sample = sample.decode()
            accepted_csv_delimiters = [',', '\t', ';']
            dialect = csv.Sniffer().sniff(sample, delimiters=accepted_csv_delimiters)
            dialect.doublequote = True  # assume that all csvs have " as string escape
        except csv.Error:
            dialect = None
        return dialect

    @staticmethod
    def _fetch_url(url: str) -> str:
        temp_dir = tempfile.mkdtemp(prefix='mindsdb_file_url_')
        try:
            r = requests.get(url, stream=True)
            if r.status_code == 200:
                with open(os.path.join(temp_dir, 'file'), 'wb') as f:
                    for chunk in r:
                        f.write(chunk)
            else:
                raise Exception(f'Responce status code is {r.status_code}')
        except Exception as e:
            print(f'Error during getting {url}')
            print(e)
            raise
        return os.path.join(temp_dir, 'file')

    def get_tables(self) -> Response:
        """
        List all files
        """
        files_meta = self.file_controller.get_files()
        data = [{
            'TABLE_NAME': x['name'],
            'TABLE_ROWS': x['row_count'],
            'TABLE_TYPE': 'BASE TABLE'
        } for x in files_meta]
        return Response(
            RESPONSE_TYPE.TABLE,
            data_frame=pd.DataFrame(data)
        )

    def get_columns(self, table_name) -> Response:
        file_meta = self.file_controller.get_file_meta(table_name)
        result = Response(
            RESPONSE_TYPE.TABLE,
            data_frame=pd.DataFrame([
                {
                    'Field': x['name'].strip() if isinstance(x, dict) else x.strip(),
                    'Type': 'str'
                } for x in file_meta['columns']
            ])
        )
        return result
