# storage.py - Enhanced for cloud storage with pre-signed URLs
from abc import ABC, abstractmethod
from pathlib import Path
from typing import BinaryIO, Union
import os
import re
from urllib.parse import urlparse
import asyncio
import aiofiles
from fastapi import UploadFile
from botocore.exceptions import ClientError
import logging

# Import global CSV utilities for better file format support
from ..preprocessing.rules.csv_utils import read_any_csv

log = logging.getLogger(__name__)


# Canonical local uploads directory (dev/local)
UPLOADS_DIR = Path(os.getenv("FM_UPLOAD_DIR", Path.cwd() / "uploads")).resolve()


def resolve_uploads_base() -> str:
    """Get the base uploads directory, same logic everywhere."""
    base = (
        os.environ.get("LOCAL_STORAGE_PATH")
        or os.environ.get("FM_UPLOAD_DIR")
        or os.path.join(os.getcwd(), "uploads")
    )
    base = os.path.abspath(base)
    # If path doesn't exist, try base/uploads (covers repo-root configs)
    if not os.path.exists(base):
        with_uploads = os.path.join(base, "uploads")
        if os.path.exists(with_uploads):
            base = with_uploads
    return base


def resolve_object_path(object_id: str) -> str:
    """
    Convert an object_id like 'uploads/<uuid>.csv' into an absolute file path
    that matches what /uploads PUT wrote and what worker_v2 reads.
    """
    base = resolve_uploads_base()
    rel = object_id.replace("\\", "/")
    if rel.startswith("uploads/"):
        rel = rel[len("uploads/") :]
    # basic traversal guard
    if ".." in rel.split("/"):
        raise ValueError("Invalid object path")
    return os.path.join(base, rel)


def exists_object(object_id: str) -> bool:
    """Check if an object exists in storage."""
    return os.path.exists(resolve_object_path(object_id))


def resolve_storage_key(key: str) -> Path:
    """Resolve a storage key or path to a concrete Path.

    Accepts:
    - Bare filenames or relative keys under FM_UPLOAD_DIR
    - Keys with accidental 'uploads/' or 'uploads\\' prefixes
    - Absolute paths
    - file:// URIs (dev only)

    Blocks path traversal outside FM_UPLOAD_DIR when resolving relative keys.
    """
    key = (key or "").strip()
    if not key:
        raise ValueError("Empty storage key")

    # file:// URI
    if key.lower().startswith("file:"):
        u = urlparse(key)
        p = Path(u.path or u.netloc).resolve()
        return p

    p = Path(key)
    # Absolute path
    if p.is_absolute():
        return p.resolve()

    # Strip accidental 'uploads/' prefix
    key = re.sub(r"^(uploads[/\\]+)", "", key)

    # Resolve under FM_UPLOAD_DIR and block traversal
    resolved = (UPLOADS_DIR / key).resolve()
    if not str(resolved).startswith(str(UPLOADS_DIR)):
        raise ValueError("Invalid storage key path traversal")
    if not resolved.exists():
        raise FileNotFoundError(
            f"Storage key not found: {key} (looked in {UPLOADS_DIR})"
        )
    return resolved


class StorageBackend(ABC):
    """Abstract storage interface"""

    @abstractmethod
    async def save(self, file: Union[UploadFile, BinaryIO], key: str) -> str:
        """Save file and return its URL/path"""
        pass

    @abstractmethod
    async def read(self, key: str) -> bytes:
        """Read file content"""
        pass

    @abstractmethod
    async def delete(self, key: str) -> bool:
        """Delete file"""
        pass

    @abstractmethod
    async def exists(self, key: str) -> bool:
        """Check if file exists"""
        pass

    @abstractmethod
    async def get_file_size(self, key: str) -> int:
        """Get file size in bytes"""
        pass

    @abstractmethod
    async def read_sample(self, key: str, max_rows: int = 1000):
        """Read sample of data as pandas DataFrame"""
        pass

    @abstractmethod
    async def generate_presigned_url(
        self, key: str, expires_in: int = 3600, operation: str = "get_object"
    ) -> str:
        """Generate a pre-signed URL for object access"""
        pass

    @abstractmethod
    async def generate_upload_url(
        self, key: str, expires_in: int = 3600, content_type: str = None
    ) -> dict:
        """Generate pre-signed URL for direct uploads"""
        pass


class LocalStorageBackend(StorageBackend):
    """Local filesystem storage (dev/testing)"""

    def __init__(self, base_path: Path):
        self.base_path = Path(base_path).expanduser().resolve()
        self.base_path.mkdir(parents=True, exist_ok=True)

    def _resolve_path(self, key: str) -> Path:
        key = (key or "").strip()
        if not key:
            raise ValueError("Empty storage key")

        log.info(
            f"[STORAGE._resolve_path] Input key: {key}, base_path: {self.base_path}"
        )

        candidate = Path(key)
        if candidate.is_absolute():
            resolved = candidate.resolve()
            try:
                resolved.relative_to(self.base_path)
            except ValueError:
                raise ValueError("Absolute storage key outside base path")
            log.info(f"[STORAGE._resolve_path] Absolute path resolved: {resolved}")
            return resolved

        normalized = key.replace("\\", "/")
        if normalized.startswith("uploads/"):
            normalized = normalized[len("uploads/") :]
        normalized = normalized.lstrip("/")

        log.info(f"[STORAGE._resolve_path] Normalized key: {normalized}")

        resolved = (self.base_path / normalized).resolve()
        try:
            resolved.relative_to(self.base_path)
        except ValueError:
            raise ValueError("Invalid storage key path traversal")

        log.info(f"[STORAGE._resolve_path] Final resolved path: {resolved}")
        return resolved

    async def save(self, file: UploadFile, key: str) -> str:
        file_path = self._resolve_path(key)
        file_path.parent.mkdir(parents=True, exist_ok=True)

        async with aiofiles.open(file_path, "wb") as f:
            content = await file.read()
            await f.write(content)

        return str(file_path)

    async def read(self, key: str) -> bytes:
        file_path = self._resolve_path(key)
        async with aiofiles.open(file_path, "rb") as f:
            return await f.read()

    async def delete(self, key: str) -> bool:
        file_path = self._resolve_path(key)
        if file_path.exists():
            file_path.unlink()
            return True
        return False

    async def exists(self, key: str) -> bool:
        try:
            path = self._resolve_path(key)
            log.info(
                f"[STORAGE.exists] Checking existence - key: {key}, resolved path: {path}, exists: {path.exists()}"
            )
        except ValueError as e:
            log.error(f"[STORAGE.exists] ValueError for key {key}: {e}")
            return False
        return path.exists()

    async def get_file_size(self, key: str) -> int:
        file_path = self._resolve_path(key)
        if file_path.exists():
            return file_path.stat().st_size
        return 0

    async def read_sample(self, key: str, max_rows: int = 800, timeout_s: int = 8):
        """
        Non-blocking sample reader for CSV/XLS/XLSX stored under base_path/key.
        - Runs pandas I/O in a worker thread
        - Caps work with an asyncio timeout
        """
        import pandas as pd

        path = self._resolve_path(key)

        async def _read_csv():
            # Prefer our csv_utils (delimiter sniffing etc.)
            return await asyncio.to_thread(read_any_csv, path)

        async def _read_xls():
            # Use kwargs to properly pass nrows parameter
            return await asyncio.to_thread(pd.read_excel, path, nrows=max_rows)

        if not path.exists():
            raise FileNotFoundError(f"Object not found: {path}")

        if path.suffix.lower() == ".csv":
            df = await asyncio.wait_for(_read_csv(), timeout=timeout_s)
            return df.head(max_rows) if max_rows else df
        elif path.suffix.lower() in [".xlsx", ".xls"]:
            # Excel files need more time to parse
            xls_timeout = min(
                timeout_s + 2, 10
            )  # Give Excel a bit more time (up to 10s)
            df = await asyncio.wait_for(_read_xls(), timeout=xls_timeout)
            return df
        else:
            raise ValueError(f"Unsupported file type: {path.suffix}")

    async def generate_presigned_url(
        self, key: str, expires_in: int = 3600, operation: str = "get_object"
    ) -> str:
        """For local storage, return a direct download endpoint"""
        # In development, return a direct API endpoint
        from urllib.parse import quote
        import time
        import json
        import base64

        # Create a simple token for local development
        token_data = {
            "key": key,
            "expires": int(time.time()) + expires_in,
            "operation": operation,
        }
        token = base64.urlsafe_b64encode(json.dumps(token_data).encode()).decode()

        # Return local API endpoint with token
        return f"/api/v1/storage/download/{quote(key)}?token={token}"

    async def generate_upload_url(
        self, key: str, expires_in: int = 3600, content_type: str = None
    ) -> dict:
        """For local storage, return upload endpoint"""
        import time
        import json
        import base64

        # Create upload token
        token_data = {
            "key": key,
            "expires": int(time.time()) + expires_in,
            "content_type": content_type or "application/octet-stream",
        }
        token = base64.urlsafe_b64encode(json.dumps(token_data).encode()).decode()

        return {
            "url": "/api/v1/storage/upload",
            "fields": {
                "key": key,
                "token": token,
                "content_type": content_type or "application/octet-stream",
            },
        }


class S3StorageBackend(StorageBackend):
    """AWS S3 storage backend"""

    def __init__(self, bucket_name: str, region: str = "us-east-1"):
        import aioboto3

        self.bucket_name = bucket_name
        self.region = region
        self._session = aioboto3.Session()

    async def save(self, file: UploadFile, key: str) -> str:
        async with self._session.client("s3", region_name=self.region) as s3:
            content = await file.read()
            await s3.put_object(
                Bucket=self.bucket_name,
                Key=key,
                Body=content,
                ContentType=file.content_type,
            )

        return f"s3://{self.bucket_name}/{key}"

    async def read(self, key: str) -> bytes:
        async with self._session.client("s3", region_name=self.region) as s3:
            response = await s3.get_object(Bucket=self.bucket_name, Key=key)
            return await response["Body"].read()

    async def delete(self, key: str) -> bool:
        async with self._session.client("s3", region_name=self.region) as s3:
            await s3.delete_object(Bucket=self.bucket_name, Key=key)
            return True

    async def exists(self, key: str) -> bool:
        async with self._session.client("s3", region_name=self.region) as s3:
            try:
                await s3.head_object(Bucket=self.bucket_name, Key=key)
                return True
            except:
                return False

    async def get_file_size(self, key: str) -> int:
        async with self._session.client("s3", region_name=self.region) as s3:
            try:
                response = await s3.head_object(Bucket=self.bucket_name, Key=key)
                return response["ContentLength"]
            except:
                return 0

    async def read_sample(self, key: str, max_rows: int = 1000):
        import pandas as pd
        import io

        # Read file content
        content = await self.read(key)

        # Determine file type from key extension
        if key.lower().endswith(".csv"):
            # Save to temp file and use global CSV reader for better format support
            import tempfile

            with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as tmp:
                tmp.write(content)
                tmp.flush()
                try:
                    df = read_any_csv(tmp.name)
                    return df.head(max_rows) if max_rows else df
                finally:
                    Path(tmp.name).unlink(missing_ok=True)
        elif key.lower().endswith((".xlsx", ".xls")):
            # Use kwargs to properly pass nrows parameter
            return pd.read_excel(io.BytesIO(content), nrows=max_rows)
        else:
            raise ValueError(f"Unsupported file type for key: {key}")

    async def generate_presigned_url(
        self, key: str, expires_in: int = 3600, operation: str = "get_object"
    ) -> str:
        """Generate a pre-signed URL for S3 object access"""
        async with self._session.client("s3", region_name=self.region) as s3:
            try:
                # Map operation to boto3 client method
                client_method = operation
                if operation == "get_object":
                    client_method = "get_object"
                elif operation == "put_object":
                    client_method = "put_object"

                url = await s3.generate_presigned_url(
                    ClientMethod=client_method,
                    Params={"Bucket": self.bucket_name, "Key": key},
                    ExpiresIn=expires_in,
                )
                return url
            except ClientError as e:
                log.error(f"Failed to generate presigned URL for {key}: {e}")
                raise ValueError(f"Could not generate presigned URL: {e}")

    async def generate_upload_url(
        self, key: str, expires_in: int = 3600, content_type: str = None
    ) -> dict:
        """Generate pre-signed POST URL for direct uploads to S3"""
        async with self._session.client("s3", region_name=self.region) as s3:
            try:
                conditions = [
                    {"bucket": self.bucket_name},
                    [
                        "starts-with",
                        "$key",
                        key.rsplit("/", 1)[0] if "/" in key else "",
                    ],
                ]
                fields = {"key": key}

                # Add content type restrictions if specified
                if content_type:
                    conditions.append({"Content-Type": content_type})
                    fields["Content-Type"] = content_type
                else:
                    # Allow common file types
                    conditions.append(["starts-with", "$Content-Type", ""])

                # Add file size limit (100MB max)
                conditions.append(["content-length-range", 0, 104857600])

                response = await s3.generate_presigned_post(
                    Bucket=self.bucket_name,
                    Key=key,
                    Fields=fields,
                    Conditions=conditions,
                    ExpiresIn=expires_in,
                )

                # Add metadata to help frontend
                response["max_file_size"] = 104857600  # 100MB
                response["expires_in"] = expires_in

                return response
            except ClientError as e:
                log.error(f"Failed to generate upload URL for {key}: {e}")
                raise ValueError(f"Could not generate upload URL: {e}")


# Factory function
def get_storage_backend() -> StorageBackend:
    """Get configured storage backend"""
    from . import settings

    if settings.STORAGE_BACKEND == "s3":
        return S3StorageBackend(
            bucket_name=settings.S3_BUCKET_NAME, region=settings.AWS_REGION
        )
    else:
        return LocalStorageBackend(settings.LOCAL_STORAGE_PATH)
