import json
import io
import boto3
import pandas as pd
from datetime import datetime, timedelta
from logger_setup import setup_logger

logger = setup_logger("recall_sampler/data_writer")

class DataWriter:
    def __init__(self, output_config: dict, customer_name: str):
        self.output_config = output_config
        self.customer_name = customer_name
        self.bucket = output_config.get("bucket")
        self.prefix = output_config.get("prefix").rstrip("/")
        self.file_format = output_config.get("file_format", "csv").lower()
        self.s3_client = boto3.client('s3')

        self.audit_log_key = f"{self.prefix}/customer={self.customer_name}/sample_audit_log.csv"

    def load_audit_log(self) -> pd.DataFrame:
        """Load audit log sample_audit_log.csv from S3 if exists."""
        try:
            response = self.s3_client.get_object(Bucket=self.bucket, Key=self.audit_log_key)
            audit_log = pd.read_csv(io.BytesIO(response['Body'].read()))
            audit_log['sampled_on'] = pd.to_datetime(audit_log['sampled_on'])
            return audit_log
        except self.s3_client.exceptions.NoSuchKey:
            # No audit log exists yet
            return pd.DataFrame(columns=["base_sku", "sampled_on", "sampling_strategy_type"])
        except Exception as e:
            raise RuntimeError(f"Failed to load audit log: {str(e)}")

    def exclude_already_sampled(self, candidates_df: pd.DataFrame, audit_log_df: pd.DataFrame) -> pd.DataFrame:
        """Exclude base_skus already sampled within last 1 year."""
        if audit_log_df.empty:
            return candidates_df

        one_year_ago = datetime.now() - timedelta(days=365)
        recent_samples = audit_log_df[audit_log_df['sampled_on'] >= one_year_ago]
        sampled_skus = set(recent_samples['base_sku'].tolist())

        filtered_df = candidates_df[~candidates_df['base_sku'].isin(sampled_skus)]
        return filtered_df

    def write_samples(self, sampled_df: pd.DataFrame, metadata: dict):
        """Write sampled data and metadata to S3 partitioned by customer/year/month."""
        now = datetime.now()
        year = now.strftime("%Y")
        month = now.strftime("%m")

        output_prefix = f"{self.prefix}/customer={self.customer_name}/year={year}/month={month}"

        # 1. Write sampled data
        sample_key = f"{output_prefix}/sampled_recall_audit.{self.file_format}"
        if self.file_format == "csv":
            csv_buffer = io.StringIO()
            sampled_df.to_csv(csv_buffer, index=False)
            self.s3_client.put_object(Bucket=self.bucket, Key=sample_key, Body=csv_buffer.getvalue())
        elif self.file_format == "parquet":
            parquet_buffer = io.BytesIO()
            sampled_df.to_parquet(parquet_buffer, index=False)
            self.s3_client.put_object(Bucket=self.bucket, Key=sample_key, Body=parquet_buffer.getvalue())
        else:
            raise ValueError(f"Unsupported file format: {self.file_format}")

        # 2. Write metadata.json
        metadata_key = f"{output_prefix}/metadata.json"
        metadata_bytes = json.dumps(metadata, indent=4).encode('utf-8')
        self.s3_client.put_object(Bucket=self.bucket, Key=metadata_key, Body=metadata_bytes)

        print(f"[SUCCESS] Sampled data written to S3: s3://{self.bucket}/{sample_key}")
        print(f"[SUCCESS] Metadata written to S3: s3://{self.bucket}/{metadata_key}")

    def update_audit_log(self, sampled_skus: list, sampling_strategy_type: str):
        """Update or create sample_audit_log.csv in S3."""
        now = datetime.now().strftime("%Y-%m-%d")

        new_entries = pd.DataFrame({
            "base_sku": sampled_skus,
            "sampled_on": now,
            "sampling_strategy_type": sampling_strategy_type
        })

        try:
            # Load existing
            response = self.s3_client.get_object(Bucket=self.bucket, Key=self.audit_log_key)
            existing_audit_log = pd.read_csv(io.BytesIO(response['Body'].read()))
            updated_audit_log = pd.concat([existing_audit_log, new_entries], ignore_index=True)
        except self.s3_client.exceptions.NoSuchKey:
            # No existing audit log, create new
            updated_audit_log = new_entries
        except Exception as e:
            raise RuntimeError(f"Failed to read existing audit log for update: {str(e)}")

        # Save updated audit log
        csv_buffer = io.StringIO()
        updated_audit_log.to_csv(csv_buffer, index=False)
        self.s3_client.put_object(Bucket=self.bucket, Key=self.audit_log_key, Body=csv_buffer.getvalue())

        print(f"[SUCCESS] Audit log updated at S3: s3://{self.bucket}/{self.audit_log_key}")