import sqlite3
import os
import time
import tempfile
import pandas as pd
import pyarrow as pa
from flask import Flask, request, jsonify
from .objstore import S3, GCS, AzureBlob
from deltalake import write_deltalake
from urllib.parse import urlparse
import logging
import logging.handlers
import threading

# Configure root logger to include timestamp
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s][PID:%(process)d][TID:%(thread)d] - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    handlers=[
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Semaphore to control concurrent writes to the additional Delta Lake path
semaphore = threading.Semaphore(1)  # Allow only one write at a time

def write_log_stats_to_deltalake(temp_file_path, org, project, cluster, deltalake_s3_path, date_int, hour_int, sqlite_s3_path):
    conn = sqlite3.connect(temp_file_path)
    query = "SELECT component, level, position, msg_digest, count, sample, sample_object_key FROM log_stats;"
    df = pd.read_sql_query(query, conn)
    conn.close()
    logger.info(f"[file={sqlite_s3_path}] Successfully read {len(df)} rows from SQLite. Transforming data...")

    if df.empty:
        return

    df['date'] = date_int
    df['hour'] = hour_int
    df['org'] = org
    df['project'] = project
    df['cluster'] = cluster

    position_split = df['position'].str.split(':', n=1, expand=True)
    position_split = position_split.reindex(columns=[0, 1], fill_value="")
    df['file'] = position_split[0].fillna('')
    df['line'] = position_split[1].fillna('')

    transformed_df = df[[
        'hour', 'date', 'org', 'project', 'cluster', 'component', 'level',
        'file', 'line', 'count', 'msg_digest', 'sample'
    ]].copy()

    transformed_df['count'] = transformed_df['count'].astype(str)
    transformed_df['file'] = transformed_df['file'].astype(str)
    transformed_df['line'] = transformed_df['line'].astype(str)
    transformed_df['level'] = transformed_df['level'].astype(str)
    transformed_df['component'] = transformed_df['component'].astype(str)
    transformed_df['msg_digest'] = transformed_df['msg_digest'].astype(str)
    transformed_df['sample'] = transformed_df['sample'].astype(str)

    delta_schema = pa.schema([
        ('hour', pa.int64()),
        ('date', pa.int64()),
        ('org', pa.string()),
        ('project', pa.string()),
        ('cluster', pa.string()),
        ('level', pa.string()),
        ('file', pa.string()),
        ('line', pa.string()),
        ('count', pa.string()),
        ('component', pa.string()),
        ('msg_digest', pa.string()),
        ('sample', pa.string()),
    ])

    pa_table = pa.Table.from_pandas(transformed_df, schema=delta_schema, preserve_index=False)
    logger.info(f"[file={sqlite_s3_path}] Transformed data to PyArrow table with {len(pa_table)} rows.")

    partition_cols = ['date', 'org', 'cluster']

    logger.info(f"[file={sqlite_s3_path}] wrote to Delta Lake path {deltalake_s3_path}")
    start_time = time.time()
    semaphore.acquire()
    start_write_time = time.time()
    success = False
    try:
        write_deltalake(
            deltalake_s3_path,
            pa_table,
            partition_by=partition_cols,
            mode='append'
        )
        success = True
    except Exception as e:
        logger.error(f"[file={sqlite_s3_path}] Error writing to additional Delta Lake: {e}")
    finally:
        semaphore.release()
    end_write_time = time.time()
    logger.info(f"[file={sqlite_s3_path}] Data processed done with success={success}, writeTime={end_write_time - start_write_time:.2f}s, blockTime={start_write_time - start_time:.2f}s")


def write_telemetry_to_deltalake(temp_file_path, org, project, cluster, deltalake_s3_path, date_int, sqlite_s3_path):
    conn = sqlite3.connect(temp_file_path)
    query = "SELECT ts, instance, telemetry_data FROM telemetry_stats;"
    df = pd.read_sql_query(query, conn)
    conn.close()
    logger.info(f"telemetry data in sqlite file [file={sqlite_s3_path}] successfully read {len(df)} rows. Transforming data...")

    if df.empty:
        return

    df['date'] = date_int
    df['org'] = org
    df['project'] = project
    df['cluster'] = cluster

    transformed_df = df[[
        'date', 'org', 'project', 'cluster', 'ts', 'instance', 'telemetry_data'
    ]].copy()

    transformed_df['ts'] = transformed_df['ts'].astype(int)
    transformed_df['instance'] = transformed_df['instance'].astype(str)
    transformed_df['telemetry_data'] = transformed_df['telemetry_data'].astype(str)

    delta_schema = pa.schema([
        ('date', pa.int64()),
        ('org', pa.string()),
        ('project', pa.string()),
        ('cluster', pa.string()),
        ('ts', pa.int64()),
        ('instance', pa.string()),
        ('telemetry_data', pa.string()),
    ])

    pa_table = pa.Table.from_pandas(transformed_df, schema=delta_schema, preserve_index=False)
    logger.info(f"telemetry data in sqlite file [file={sqlite_s3_path}] transformed to PyArrow table with {len(pa_table)} rows.")

    partition_cols = ['date', 'org', 'cluster']

    logger.info(f"telemetry data in sqlite file [file={sqlite_s3_path}] wrote to Delta Lake path {deltalake_s3_path}")
    start_time = time.time()
    semaphore.acquire()
    start_write_time = time.time()
    success = False
    try:
        write_deltalake(
            deltalake_s3_path,
            pa_table,
            partition_by=partition_cols,
            mode='append'
        )
        success = True
    except Exception as e:
        logger.error(f"telemetry data in sqlite file [file={sqlite_s3_path}] Error writing to additional Delta Lake: {e}")
    finally:
        semaphore.release()
    end_write_time = time.time()
    logger.info(f"telemetry data in sqlite file [file={sqlite_s3_path}] Data processed done with success={success}, writeTime={end_write_time - start_write_time:.2f}s, blockTime={start_write_time - start_time:.2f}s")

def process_sqlite_to_deltalake(
        vendor: str,
        sqlite_s3_path: str,
        org: str,
        project: str,
        cluster: str,
        deltalake_s3_base_path: str
):
    """
    Downloads a SQLite file from S3, processes its log_stats data, and writes
    it to a Delta Lake table on S3, partitioned by date, org, and cluster.
    Supports writing to an additional Delta Lake path if provided.

    Args:
        vendor: The vendor name. Supported values: "aws", "azure", "gcp" and "alicloud".
        sqlite_s3_path: The full S3 path to the SQLite file
                        (e.g., "s3://bucket/path/to/2025051519.db").
        org: The organization name.
        project: The project name.
        cluster: The cluster name.
        deltalake_s3_base_path: A base S3 path for writing Delta Lake data, including log-stats and telemetry
    """
    logger.info(f"[file={sqlite_s3_path}] Starting processing request")

    try:
        parsed_url = urlparse(sqlite_s3_path)
        if parsed_url.scheme not in ['s3', 'gs', 'az']:
            raise ValueError(f"Invalid S3 path scheme: {parsed_url.scheme}")
        s3_bucket = parsed_url.netloc
        s3_key = parsed_url.path.lstrip('/')
        file_name = os.path.basename(s3_key)

        if len(file_name) < 11 or not file_name.endswith('.db'):
            raise ValueError(f"Filename format unexpected: {file_name}")
        date_hour_str = file_name[:-3]
        if len(date_hour_str) != 10 or not date_hour_str.isdigit():
            raise ValueError(f"Date/hour string format unexpected: {date_hour_str}")

        date_str = date_hour_str[:8]
        hour_str = date_hour_str[8:]

        date_int = int(date_str)
        hour_int = int(hour_str)
        logger.info(f"[file={sqlite_s3_path}] Parsed: bucket={s3_bucket}, key={s3_key}, date={date_int}, hour={hour_int}")

    except Exception as e:
        logger.error(f"[file={sqlite_s3_path}] Error processing request: {str(e)}")
        return

    # If you want to use alicloud with the same code as aws, you should set the env like
    # https://github.com/delta-io/delta-rs/issues/2361#issuecomment-2176609082
    if vendor =='aws' or vendor == "alicloud":
        s3_client = S3()
    elif vendor == 'gcp':
        s3_client = GCS()
    elif vendor == 'azure':
        s3_client = AzureBlob()
    else:
        raise ValueError(f"Invalid vendor: {vendor}")

    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as fp:
            temp_file_path = fp.name
        logger.info(f"[file={sqlite_s3_path}] Downloading s3 file to temporary file {temp_file_path}")
        succ = s3_client.download(s3_bucket, s3_key, temp_file_path)
        if not succ:
            raise FileNotFoundError(f"Download failed for file {s3_bucket}/{s3_key}")
        logger.info(f"[file={sqlite_s3_path}] Download complete, Reading data from SQLite file: {temp_file_path}")

        write_log_stats_to_deltalake(temp_file_path, org, project, cluster, deltalake_s3_base_path+"log-stats", date_int, hour_int, sqlite_s3_path)
        write_telemetry_to_deltalake(temp_file_path, org, project, cluster, deltalake_s3_base_path+"telemetry", date_int, sqlite_s3_path)

    except FileNotFoundError as e:
        logger.error(f"[file={sqlite_s3_path}] Error: Download failed, temporary file not found. {e}")
    except Exception as e:
        logger.error(f"[file={sqlite_s3_path}] An error occurred during processing: {e}")

    finally:
        if 'temp_file_path' in locals() and os.path.exists(temp_file_path):
            try:
                os.remove(temp_file_path)
                logger.info(f"[file={sqlite_s3_path}] Cleaned up temporary file: {temp_file_path}")
            except Exception as e:
                logger.error(f"[file={sqlite_s3_path}] Error cleaning up temporary file {temp_file_path}: {e}")

def write_to_deltalake(vendor, org, project, cluster, bucket, sqlite_s3_path):
    if vendor == "aws" or vendor == "alicloud":
        deltalake_base_path = f"s3://{bucket}/deltalake/global/"
    elif vendor == "gcp":
        deltalake_base_path = f"gs://{bucket}/deltalake/global/"
    elif vendor == "azure":
        deltalake_base_path = f"abfs://{bucket}/deltalake/global/"
    else:
        raise ValueError(f"Invalid vendor: {vendor}")

    process_sqlite_to_deltalake(
        vendor=vendor,
        sqlite_s3_path=sqlite_s3_path,
        org=org,
        project=project,
        cluster=cluster,
        deltalake_s3_base_path=deltalake_base_path
    )

@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({"message": "healthy"}), 200


@app.route('/deltalake/log_stats', methods=['POST'])
def write_log_state():
    data = request.get_json()
    if not data:
        return jsonify({"error": "need data"}), 400
    logger.info(f"Received request data: {data}")
    if data['vendor'] not in ["aws", "azure", "gcp", "alicloud"]:
        return jsonify({"error": "invalid vendor"}), 400

    write_to_deltalake(data['vendor'], data['org'], data['project'], data['cluster'], data['bucket'], data['sqlite_s3_path'])
    return jsonify({"message": "task accepted"}), 202

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8080)