import gen3_validator
import os
import tempfile
import json
import boto3
import pandas as pd
import logging
import awswrangler as wr
import io
from botocore.exceptions import ClientError

logger = logging.getLogger(__name__)

def parse_s3_uri(s3_uri: str):
    """Parse an S3 URI into (bucket, key)"""
    if not s3_uri.startswith("s3://"):
        logger.error("s3_uri must start with 's3://'")
        raise ValueError("s3_uri must start with 's3://'")
    _, _, bucket_and_key = s3_uri.partition("s3://")
    bucket, _, key = bucket_and_key.partition("/")
    if not bucket or not key:
        logger.error("Malformed s3_uri: must be of the form 's3://bucket/key'")
        raise ValueError("Malformed s3_uri: must be of the form 's3://bucket/key'")
    return bucket, key

def get_s3_client():
    """Return a cached boto3 S3 client instance."""
    # Could cache if perf-critical, otherwise return new for now
    try:
        return boto3.client('s3')
    except Exception as e:
        logger.error(f"Failed to create boto3 S3 client: {e}")
        raise

def read_json_from_s3(s3_uri: str):
    try:
        bucket, key = parse_s3_uri(s3_uri)
        s3 = get_s3_client()
        obj = s3.get_object(Bucket=bucket, Key=key)
        return json.load(obj['Body'])
    except ClientError as ce:
        logger.error(f"ClientError reading JSON from {s3_uri}: {ce}")
        raise
    except Exception as e:
        logger.error(f"Failed to read JSON from {s3_uri}: {e}")
        raise

def write_bytes_to_s3(bytes_data, s3_uri: str):
    try:
        bucket, key = parse_s3_uri(s3_uri)
        s3 = get_s3_client()
        s3.put_object(Bucket=bucket, Key=key, Body=bytes_data)
        logger.info(f"Successfully wrote to {s3_uri}")
    except Exception as e:
        logger.error(f"Failed to write bytes to {s3_uri}: {e}")
        raise

def download_s3_file(s3_uri: str, local_path: str):
    bucket, key = parse_s3_uri(s3_uri)
    s3 = get_s3_client()
    try:
        logger.info(f"Downloading {s3_uri} to {local_path}")
        s3.download_file(bucket, key, local_path)
        logger.debug(f"Downloaded: {local_path}")
    except ClientError as ce:
        logger.error(f"ClientError while downloading {s3_uri} to {local_path}: {ce}")
        raise
    except Exception as e:
        logger.error(f"Failed to download {s3_uri} to {local_path}: {e}")
        raise

def list_s3_objects(s3_uri: str):
    try:
        bucket, prefix = parse_s3_uri(s3_uri)
        if prefix and not prefix.endswith("/"):
            prefix += "/"
        s3 = get_s3_client()
        paginator = s3.get_paginator('list_objects_v2')
        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
            contents = page.get("Contents", [])
            for obj in contents:
                yield bucket, obj["Key"]
    except ClientError as ce:
        logger.error(f"ClientError listing objects for {s3_uri}: {ce}")
        raise
    except Exception as e:
        logger.error(f"Error listing objects for {s3_uri}: {e}")
        raise

def load_schema_from_s3_uri(s3_uri: str) -> dict:
    logger.info(f"Loading schema from S3 URI: {s3_uri}")
    try:
        schema = read_json_from_s3(s3_uri)
        logger.info(f"Schema loaded successfully from {s3_uri}")
        return schema
    except Exception as e:
        logger.error(f"Could not load schema from {s3_uri}: {e}")
        raise

def write_schema_to_temp_file(schema: dict, filename: str = "schema.json") -> str:
    try:
        temp_dir = os.path.join("/tmp", next(tempfile._get_candidate_names()))
        os.makedirs(temp_dir, exist_ok=True)
        schema_path = os.path.join(temp_dir, filename)
        logger.info(f"Writing schema to temp file: {schema_path}")
        with open(schema_path, "w", encoding="utf-8") as f:
            json.dump(schema, f, indent=2)
        logger.info(f"Schema written to: {schema_path}")
        return os.path.abspath(schema_path)
    except Exception as e:
        logger.error(f"Failed to write schema to temp file: {e}")
        raise

def create_metadata_table(s3_uri: str) -> list[dict]:
    import re
    logger.info(f"Creating metadata table from S3 URI: {s3_uri}")
    result = []
    regex = re.compile(
        r"study_id=([^/]+)/validation_id=([^/]+)/table_name=([^/]+)/snapshot_id=([^/]+)/([^/]+\.json)$"
    )
    try:
        for bucket, key in list_s3_objects(s3_uri):
            if not key.endswith(".json"):
                continue
            match = regex.search(key)
            if match:
                study_id, validation_id, table_name, snapshot_id, file_name = match.groups()
                entry = {
                    "study_id": study_id,
                    "validation_id": validation_id,
                    "table_name": table_name,
                    "snapshot_id": snapshot_id,
                    "file_name": file_name,
                    "s3_uri": f"s3://{bucket}/{key}"
                }
                logger.debug(f"Matched file entry: {entry}")
                result.append(entry)
        logger.info(f"Finished building metadata table. {len(result)} records found.")
        return result
    except Exception as e:
        logger.error(f"Error occurred while creating metadata table from {s3_uri}: {e}")
        raise

def get_latest_validation_for_study(metadata_table: pd.DataFrame, study_id: str):
    logger.info(f"Selecting latest validation for study_id: {study_id}")
    try:
        if 'study_id' not in metadata_table.columns:
            logger.warning("No 'study_id' column in metadata_table.")
            return pd.DataFrame(), None
        filtered_table = metadata_table[metadata_table['study_id'] == study_id]
        if not filtered_table.empty:
            latest_validation_id = filtered_table['validation_id'].max()
            logger.info(f"Latest validation_id for study {study_id}: {latest_validation_id}")
            result_table = filtered_table[filtered_table['validation_id'] == latest_validation_id]
        else:
            logger.warning(f"No entries found for study_id {study_id}")
            result_table = pd.DataFrame()
            latest_validation_id = None
        logger.debug(f"{len(result_table)} records returned for latest validation_id.")
        return result_table, latest_validation_id
    except Exception as e:
        logger.error(f"Error attempting to get latest validation for study {study_id}: {e}")
        raise

def download_s3_files_to_temp_dir(s3_uris: list[str], temp_dir_name: str = "to_validate"):
    logger.info(f"Downloading {len(s3_uris)} S3 files to temp dir named '{temp_dir_name}'")
    try:
        base_tempdir = tempfile.mkdtemp()
        target_temp_dir = os.path.join(base_tempdir, temp_dir_name)
        os.makedirs(target_temp_dir, exist_ok=True)
        logger.debug(f"Target temp dir for downloads: {target_temp_dir}")
        downloaded_files = []
        for s3_uri in s3_uris:
            if not s3_uri.startswith("s3://"):
                logger.warning(f"Skipping invalid s3_uri (not an s3 path): {s3_uri}")
                continue
            _, _, path = s3_uri.partition("s3://")
            file_name = os.path.basename(path.split("/", 1)[1] if "/" in path else path)
            parts = file_name.split("_")
            if len(parts) > 2:
                final_file_name = "_".join(parts[2:])
            else:
                final_file_name = file_name
            dest_path = os.path.join(target_temp_dir, final_file_name)
            try:
                download_s3_file(s3_uri, dest_path)
                downloaded_files.append(dest_path)
            except Exception as e:
                logger.warning(f"Skipping failed download for {s3_uri}: {e}")
                continue
        logger.info(f"Total files downloaded: {len(downloaded_files)}")
        return target_temp_dir, downloaded_files
    except Exception as e:
        logger.error(f"Failed to download S3 files to temp dir: {e}")
        raise

def write_df_to_s3(df: pd.DataFrame, s3_uri: str, filename: str, sep: str = ",", index: bool = False):
    logger.info(f"Writing DataFrame to S3: {s3_uri}, filename: {filename}")
    try:
        if s3_uri.endswith("/"):
            s3_uri = s3_uri[:-1]
        s3_path = f"{s3_uri}/{filename}"
        csv_buffer = io.StringIO()
        df.to_csv(csv_buffer, sep=sep, index=index)
        write_bytes_to_s3(csv_buffer.getvalue(), s3_path)
    except Exception as e:
        logger.error(f"Failed to write DataFrame to S3 ({s3_uri}, {filename}): {e}")
        raise

def write_parquet_to_db(
    df: pd.DataFrame,
    dataset_root: str,
    database: str,
    table: str,
    partition_cols: list[str] = ["study_id", "submission_date"],
    compression: str = "snappy",
    mode: str = "append",
    schema_evolution: bool = True,
) -> None:
    try:
        logger.debug(f"Creating Glue database '{database}' if not exists.")
        wr.catalog.create_database(name=database, exist_ok=True)
        logger.debug(
            f"Writing DataFrame to Parquet at {dataset_root.rstrip('/')}/ (table: {table}, database: {database}, "
            f"partitions: {partition_cols}, compression: {compression}, mode: {mode}, schema_evolution: {schema_evolution})"
        )
        wr.s3.to_parquet(
            df=df.astype("string"),
            path=dataset_root.rstrip("/") + "/",
            dataset=True,
            database=database,
            table=table,
            partition_cols=partition_cols,
            compression=compression,
            mode=mode,
            schema_evolution=schema_evolution,
        )
        logger.debug(f"Successfully wrote Parquet dataset to S3 at {dataset_root} (table: {table}, database: {database})")
    except Exception as e:
        logger.error(
            f"Failed to write Parquet dataset to S3 at {dataset_root} (table: {table}, database: {database}): {e}"
        )
        raise RuntimeError(
            f"Failed to write Parquet dataset to S3 at {dataset_root} (table: {table}, database: {database}): {e}"
        )


def validate_pipeline(
    study_id: str,
    schema_s3_uri: str,
    validation_s3_uri: str,
    write_back_root: str,
    parquet_root: str,
    glue_database: str
):
    """
    Orchestrates validation of study data against a Gen3 schema, downloading validation results,
    resolving the schema, and updating results in an AWS Glue/Athena-compatible Parquet database.

    Parameters
    ----------
    study_id : str, optional
        The identifier for the study whose validation files and results you wish to process.
        This should match the `study_id` field in your metadata and file naming conventions.
        Example: "ausdiab"
    
    schema_s3_uri : str, optional
        The S3 URI for the Gen3-compliant JSON schema to use for validation.
        This should be a full S3 path to a JSON schema file that describes what
        the study data should look like (fields, types, structure, etc).
        Example: "s3://gen3schema-cad-uat-biocommons.org.au/cad.json"
    
    validation_s3_uri : str, optional
        The S3 URI (optionally a folder/prefix) under which all validation result
        files are written or stored. This is where the function will look for
        previously validated files (e.g., json/csv result files) to find the
        most recent validation outputs for the given study.
        Example: "s3://acdc-staging-raw-silver/validation/"
    
    write_back_root : str, optional
        The S3 URI prefix/folder where updated or newly generated validation results
        (such as summary reports or processed files) should be written back.
        By default, this matches the validation_s3_uri.
        Example: "s3://acdc-staging-raw-silver/validation/"
    
    parquet_root : str, optional
        The S3 URI prefix/folder to which Parquet files compatible with AWS Athena/Glue
        should be written. These files summarize or contain results of the validation,
        and can be queried from AWS Athena.
        Example: "s3://acdc-staging-raw-silver/validation/athena/"
    
    glue_database : str, optional
        The name of the AWS Glue database to use for registering the Parquet data.
        If the database does not exist, it will be created. This allows the
        validation results to be queried via AWS Athena/Glue queries.
        Example: "acdc_validation_db"
    """

    try:
        schema = load_schema_from_s3_uri(schema_s3_uri)
        schema_path = write_schema_to_temp_file(schema)
    except Exception as e:
        logger.error(f"Schema loading or writing failed: {e}")
        raise

    try:
        logger.info(f"Instantiating schema resolver with schema path: {schema_path}")
        resolver = gen3_validator.ResolveSchema(schema_path=schema_path)
        resolver.resolve_schema()
        logger.info("Schema resolved.")
    except Exception as e:
        logger.error(f"Failed to instantiate or resolve schema: {e}")
        raise
    
    try:
        logger.info(f"Getting schema version")
        dd = gen3_validator.DataDictionary(schema_path = schema_path)
        dd.parse_schema()
        schema_version = dd.get_schema_version(dd.schema)
        logger.info(f"Schema version: {schema_version}")
    except Exception as e:
        logger.error(f"Failed to get schema version: {e}")
        raise

    try:
        metadata_table = pd.DataFrame(create_metadata_table(validation_s3_uri))
        logger.info(f"Metadata table created from S3 ({len(metadata_table)} rows).")
    except Exception as e:
        logger.error(f"Failed to create metadata table: {e}")
        raise

    try:
        latest_validation_metadata_table, latest_validation_id = get_latest_validation_for_study(metadata_table, study_id)
        if latest_validation_metadata_table is None or latest_validation_metadata_table.empty:
            logger.error(f"No latest validation files found for study {study_id}.")
            raise Exception(f"No latest validation files found for study {study_id}.")

        latest_validation_file_names = latest_validation_metadata_table['file_name'].tolist()
        latest_validation_s3_uris = latest_validation_metadata_table['s3_uri'].tolist()
        logger.info(f"Latest validation file names: {latest_validation_file_names}")
        logger.info(f"Latest validation S3 URIs: {latest_validation_s3_uris}")

        downloaded_files_dir, downloaded_files = download_s3_files_to_temp_dir(latest_validation_s3_uris)
    except Exception as e:
        logger.error(f"Failed to retrieve and download latest validation files: {e}")
        raise

    try:
        files_in_dir = os.listdir(downloaded_files_dir)
        logger.info(f"Files in downloaded_files_dir: {downloaded_files_dir}")
        for f in files_in_dir:
            logger.info(f"  - {f}")
    except Exception as e:
        logger.error(f"Failed to list or log downloaded files: {e}")
        raise

    try:
        logger.info(f"Starting validation for files in: {downloaded_files_dir}")
        data = gen3_validator.ParseData(data_folder_path=downloaded_files_dir)
        validator = gen3_validator.Validate(data_map=data.data_dict, resolved_schema=resolver.schema_resolved)
        validator.validate_schema()
        logger.info("Validation completed.")

        summary = gen3_validator.ValidateSummary(validator)
        summary.flatten_validation_results()
        full_validation_results_df = summary.flattened_results_to_pd()
        full_validation_results_df['validation_id'] = latest_validation_id
        full_validation_results_df['study_id'] = study_id
        full_validation_results_df['schema_version'] = schema_version
        summarised_validation_results_df = summary.collapse_flatten_results_to_pd()
        summarised_validation_results_df['validation_id'] = latest_validation_id
        summarised_validation_results_df['study_id'] = study_id
        summarised_validation_results_df['schema_version'] = schema_version
        logger.info("Validation results processed to DataFrame.")
    except Exception as e:
        logger.error(f"Validation failed: {e}")
        raise

    write_back_s3_uri = f"{write_back_root.rstrip('/')}/study_id={study_id}/validation_id={latest_validation_id}/"
    try:
        logger.info(f"Writing out validation results to S3 location: {write_back_s3_uri}")
        write_df_to_s3(full_validation_results_df, write_back_s3_uri, "full_validation_results.csv")
        write_df_to_s3(summarised_validation_results_df, write_back_s3_uri, "summarised_validation_results.csv")
    except Exception as e:
        logger.error(f"Failed to write CSV validation results to S3: {e}")
        raise

    try:
        logger.info(f"Writing full_validation_results to database '{glue_database}', table 'full_validation_results'")
        write_parquet_to_db(
            df=full_validation_results_df,
            dataset_root=parquet_root,
            database=glue_database,
            table="full_validation_results",
            partition_cols=["validation_id"],
            compression="snappy",
            mode="append",
            schema_evolution=True
        )
        logger.info(f"Writing summarised_validation_results to database '{glue_database}', table 'summarised_validation_results'")
        write_parquet_to_db(
            df=summarised_validation_results_df,
            dataset_root=parquet_root,
            database=glue_database,
            table="summarised_validation_results",
            partition_cols=["validation_id"],
            compression="snappy",
            mode="append",
            schema_evolution=True
        )
    except Exception as e:
        logger.error(f"Failed to write validation results to Parquet/Glue: {e}")
        raise