import os
import pandas as pd
import polars as pl
import json
import pickle as pkl
import warnings
import boto3
from botocore.exceptions import ClientError
from typing import Union, Optional
from io import BytesIO
from datetime import datetime, timedelta, timezone
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
AWS_REGION = os.getenv("AWS_REGION")


class S3CloudHelper:
    def __init__(
        self,
        obj: Union[pd.DataFrame, pl.DataFrame, dict, str, None] = None,
        path: str = None,
        region_name: str = AWS_REGION,
    ):
        if obj is not None and path is not None:
            raise ValueError(
                "Only one of 'obj' or 'path' should be provided, not both."
            )

        self.obj = obj
        self.path = path

        self.s3_client = boto3.client(
            "s3",
            region_name=region_name,
            aws_access_key_id=AWS_ACCESS_KEY_ID,
            aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
        )

    def _infer_file_type(self, file_name: str) -> Optional[str]:
        ext = os.path.splitext(file_name)[1].lower()
        return {
            ".csv": "csv",
            ".json": "json",
            ".pkl": "pickle",
            ".pickle": "pickle",
            ".parquet": "parquet",
            ".pq": "parquet",
            ".txt": "txt",
            "": None,
        }.get(ext, None)

    def upload_to_s3(
        self,
        bucket_name: str,
        file_name: str,
        file_type: str = None,
    ):
        if file_type is None:
            file_type = self._infer_file_type(file_name)
            if file_type is None:
                raise ValueError(f"Cannot infer file_type from filename: {file_name}")

        if self.path:
            self.s3_client.upload_file(self.path, bucket_name, file_name)
            return

        if self.obj is not None:
            buffer = BytesIO()

            if file_type == "csv":
                if isinstance(self.obj, pd.DataFrame):
                    self.obj.to_csv(buffer, index=False)
                elif isinstance(self.obj, pl.DataFrame):
                    buffer.write(self.obj.write_csv().encode("utf-8"))
                buffer.seek(0)
                self.s3_client.upload_fileobj(buffer, bucket_name, file_name)

            elif file_type == "parquet":
                if isinstance(self.obj, pd.DataFrame):
                    self.obj.to_parquet(buffer, index=False)
                elif isinstance(self.obj, pl.DataFrame):
                    self.obj.write_parquet(buffer)
                buffer.seek(0)
                self.s3_client.upload_fileobj(buffer, bucket_name, file_name)

            elif file_type == "json":
                json_str = json.dumps(self.obj, indent=2)
                self.s3_client.put_object(
                    Body=json_str.encode("utf-8"), Bucket=bucket_name, Key=file_name
                )

            elif file_type in ("pickle", "pkl"):
                pkl.dump(self.obj, buffer)
                buffer.seek(0)
                self.s3_client.upload_fileobj(buffer, bucket_name, file_name)

            elif file_type == "txt":
                self.s3_client.put_object(
                    Body=self.obj.encode("utf-8"), Bucket=bucket_name, Key=file_name
                )
            else:
                raise ValueError(f"Unsupported file type {file_type}")

    def download_from_s3(
        self,
        s3_filepath: str,
        file_type: str = None,
        use_polars: bool = False,
    ):
        if s3_filepath.startswith("s3://"):
            s3_filepath = s3_filepath[5:]

        bucket_name, *blob_path = s3_filepath.split("/", 1)
        blob_path = blob_path[0]

        is_prefix = not any(
            blob_path.endswith(ext)
            for ext in [".csv", ".parquet", ".json", ".pkl", ".pickle", ".txt"]
        )

        if is_prefix:
            files = self.list_files(bucket=bucket_name, prefix=blob_path)
            if not files:
                return pl.DataFrame() if use_polars else pd.DataFrame()

            if use_polars:
                if file_type == "csv":
                    return pl.scan_csv(files)
                elif file_type == "parquet":
                    return pl.scan_parquet(files)
            else:
                raise ValueError("Lazy multi-file read only supported with Polars.")
        else:
            response = self.s3_client.get_object(Bucket=bucket_name, Key=blob_path)
            data = response["Body"].read()
            buffer = BytesIO(data)

            if file_type is None:
                file_type = self._infer_file_type(blob_path)

            if file_type == "csv":
                return pl.read_csv(buffer) if use_polars else pd.read_csv(buffer)
            elif file_type == "parquet":
                return (
                    pl.read_parquet(buffer) if use_polars else pd.read_parquet(buffer)
                )
            elif file_type == "json":
                return json.loads(data.decode("utf-8"))
            elif file_type in ("pickle", "pkl"):
                return pkl.loads(data)
            elif file_type == "txt":
                return data.decode("utf-8")

    def delete_from_s3(self, bucket_name: str, file_name: str) -> bool:
        try:
            if file_name.endswith("/"):
                paginator = self.s3_client.get_paginator("list_objects_v2")
                pages = paginator.paginate(Bucket=bucket_name, Prefix=file_name)
                objects_to_delete = [
                    {"Key": obj["Key"]}
                    for page in pages
                    for obj in page.get("Contents", [])
                ]
                if objects_to_delete:
                    for i in range(0, len(objects_to_delete), 1000):
                        batch = objects_to_delete[i : i + 1000]
                        self.s3_client.delete_objects(
                            Bucket=bucket_name, Delete={"Objects": batch, "Quiet": True}
                        )
                    return True
                return False
            else:
                self.s3_client.delete_object(Bucket=bucket_name, Key=file_name)
                return True
        except ClientError as e:
            warnings.warn(f"Failed to delete: {file_name}: {e}")
            return False

    def list_files(self, bucket: str, prefix: str) -> list[str]:
        try:
            response = self.s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix)
            return [
                f"s3://{bucket}/{obj['Key']}"
                for obj in response.get("Contents", [])
                if not obj["Key"].endswith("/")
            ]
        except ClientError as e:
            warnings.warn(
                f"Failed to list files in bucket '{bucket}' with prefix '{prefix}': {e}"
            )
            return []

    def update_and_archive_table(
        self,
        bucket_name: str,
        table_name: str,
        file_name: str,
        file_type: str = None,
        archive_retention_days: int = 30,
        upsert: bool = False,
        dedup_columns: Optional[list[str]] = None,
    ):
        """
        Rotate table data by:
        1. Moving current files to archive with timestamp (skip if upsert=True)
        2. Uploading new data to current folder (or appending and deduplicating if upsert=True)
        3. Deleting archived files older than retention_days

        Args:
            bucket_name: S3 bucket name
            table_name: Name of the table (used as folder prefix)
            file_name: Name for the new file in current folder
            file_type: File type (csv, parquet, json, etc.). If None, inferred from file_name
            archive_retention_days: Days to keep archived files (default: 30)
            upsert: If True, append to existing current file and drop duplicates (default: False)
            dedup_columns: Columns to use for dropping duplicates in upsert mode. If None, uses all columns
        """
        current_prefix = f"{table_name}/current/"
        archive_prefix = f"{table_name}/archive/"
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        # Determine file type and ensure proper extension
        if file_type is None:
            file_type = self._infer_file_type(file_name)
            if file_type is None:
                raise ValueError(f"Cannot infer file_type from filename: {file_name}")
        
        # Extract name and extension, build proper extension if needed
        name_without_ext, ext = os.path.splitext(file_name)
        if not ext:
            # No extension in filename, build it from file_type
            ext_map = {
                "csv": ".csv",
                "json": ".json",
                "pickle": ".pkl",
                "pkl": ".pkl",
                "parquet": ".parquet",
                "txt": ".txt"
            }
            ext = ext_map.get(file_type, "")
        
        timestamped_filename = f"{timestamp}_{name_without_ext}{ext}"
        new_file_key = f"{current_prefix}{timestamped_filename}"

        # If upsert mode, download existing data before archiving
        existing_df = None
        if upsert and (file_type in ("csv", "parquet")):
            if not isinstance(self.obj, (pd.DataFrame, pl.DataFrame)):
                raise ValueError("Upsert mode only supports DataFrame objects (pandas or polars)")
            
            try:
                existing_files = self.list_files(bucket=bucket_name, prefix=current_prefix)
                
                if existing_files:
                    # Download the first file found in current
                    existing_s3_path = existing_files[0]
                    use_polars = isinstance(self.obj, pl.DataFrame)
                    existing_df = self.download_from_s3(
                        s3_filepath=existing_s3_path,
                        file_type=file_type,
                        use_polars=use_polars
                    )
                    
            except Exception as e:
                warnings.warn(f"Error downloading existing file for upsert: {e}. Will upload as new file.")

        # Step 1: Move existing files from current/ to archive/
        if not upsert:
            # Step 1: Move existing files from current/ to archive/
            try:
                paginator = self.s3_client.get_paginator("list_objects_v2")
                pages = paginator.paginate(Bucket=bucket_name, Prefix=current_prefix)

                for page in pages:
                    for obj in page.get("Contents", []):
                        source_key = obj["Key"]
                        if source_key.endswith("/"):
                            continue

                        # Extract filename and extension from path
                        filename_only = os.path.basename(source_key)
                        name_without_ext, ext = os.path.splitext(filename_only)
                        # Create archived filename with timestamp and extension
                        archive_key = f"{archive_prefix}{timestamp}_{name_without_ext}{ext}"

                        # Copy to archive
                        self.s3_client.copy_object(
                            Bucket=bucket_name,
                            CopySource={"Bucket": bucket_name, "Key": source_key},
                            Key=archive_key,
                        )

                        # Delete from current
                        self.s3_client.delete_object(Bucket=bucket_name, Key=source_key)

            except ClientError as e:
                warnings.warn(f"Error moving files to archive: {e}")

        # Step 1.5: Ensure no more than 2 files per base filename from today in archive
        try:
            today_prefix = datetime.now().strftime("%Y%m%d")
            paginator = self.s3_client.get_paginator("list_objects_v2")
            pages = paginator.paginate(Bucket=bucket_name, Prefix=archive_prefix)

            # Group files by base filename (without timestamp prefix)
            files_by_base = {}
            for page in pages:
                for obj in page.get("Contents", []):
                    if obj["Key"].endswith("/"):
                        continue
                    
                    filename = os.path.basename(obj["Key"])
                    # Check if filename starts with today's date
                    if filename.startswith(today_prefix):
                        # Extract base filename (everything after timestamp_)
                        # Format is: YYYYMMDD_HHMMSS_originalname
                        parts = filename.split("_", 2)
                        if len(parts) >= 3:
                            base_name = parts[2]
                            if base_name not in files_by_base:
                                files_by_base[base_name] = []
                            files_by_base[base_name].append({
                                "key": obj["Key"],
                                "last_modified": obj["LastModified"]
                            })

            # For each base filename, keep only the 2 most recent
            for base_name, files in files_by_base.items():
                if len(files) > 2:
                    # Sort by last_modified (newest first)
                    files.sort(key=lambda x: x["last_modified"], reverse=True)
                    # Delete all but the 2 most recent
                    for file_to_delete in files[2:]:
                        self.s3_client.delete_object(
                            Bucket=bucket_name, Key=file_to_delete["key"]
                        )

        except ClientError as e:
            warnings.warn(f"Error cleaning up duplicate archived files from today: {e}")

        # Step 2: Upload new data to current/ (or upsert)
        # If upsert mode and we have existing data, merge it
        if upsert and existing_df is not None:
            use_polars = isinstance(self.obj, pl.DataFrame)
            
            # Append new data
            if use_polars:
                combined_df = pl.concat([existing_df, self.obj])
                # Drop duplicates
                if dedup_columns:
                    combined_df = combined_df.unique(subset=dedup_columns, keep="last")
                else:
                    combined_df = combined_df.unique(keep="last")
            else:
                combined_df = pd.concat([existing_df, self.obj], ignore_index=True)
                # Drop duplicates
                if dedup_columns:
                    combined_df = combined_df.drop_duplicates(subset=dedup_columns, keep="last")
                else:
                    combined_df = combined_df.drop_duplicates(keep="last")
            
            # Update self.obj to the combined dataframe
            self.obj = combined_df
        
        # Upload the data (new or combined)
        if self.path:
            buffer = BytesIO()

            if file_type == "csv":
                if isinstance(self.obj, pd.DataFrame):
                    self.obj.to_csv(buffer, index=False)
                elif isinstance(self.obj, pl.DataFrame):
                    buffer.write(self.obj.write_csv().encode("utf-8"))
                buffer.seek(0)
                self.s3_client.upload_fileobj(buffer, bucket_name, new_file_key)

            elif file_type == "parquet":
                if isinstance(self.obj, pd.DataFrame):
                    self.obj.to_parquet(buffer, index=False)
                elif isinstance(self.obj, pl.DataFrame):
                    self.obj.write_parquet(buffer)
                buffer.seek(0)
                self.s3_client.upload_fileobj(buffer, bucket_name, new_file_key)

            elif file_type == "json":
                json_str = json.dumps(self.obj, indent=2)
                self.s3_client.put_object(
                    Body=json_str.encode("utf-8"), Bucket=bucket_name, Key=new_file_key
                )

            elif file_type in ("pickle", "pkl"):
                pkl.dump(self.obj, buffer)
                buffer.seek(0)
                self.s3_client.upload_fileobj(buffer, bucket_name, new_file_key)

            elif file_type == "txt":
                self.s3_client.put_object(
                    Body=self.obj.encode("utf-8"), Bucket=bucket_name, Key=new_file_key
                )
            else:
                raise ValueError(f"Unsupported file type {file_type}")
        else:
            raise ValueError("No data to upload. Provide either 'obj' or 'path'.")

        # Step 3: Delete old archived files
        cutoff_date = datetime.now() - timedelta(days=archive_retention_days)

        try:
            paginator = self.s3_client.get_paginator("list_objects_v2")
            pages = paginator.paginate(Bucket=bucket_name, Prefix=archive_prefix)

            for page in pages:
                for obj in page.get("Contents", []):
                    if obj["Key"].endswith("/"):
                        continue

                    # Check if file is older than retention period
                    last_modified = obj["LastModified"].astimezone()
                    if last_modified.replace(tzinfo=None) < cutoff_date:
                        self.s3_client.delete_object(
                            Bucket=bucket_name, Key=obj["Key"]
                        )

        except ClientError as e:
            warnings.warn(f"Error cleaning up old archived files: {e}")
