"""
Storage backends for pasted traces.
"""
import json
import re
from pathlib import Path
from typing import Dict, List, Optional


class BaseTracesFolderService:
    """
    Backend abstraction for pasted trace persistence.
    """

    def __init__(self, logger, storage_prefix: str):
        self.logger = logger
        self.storage_prefix = storage_prefix

    def _sanitize_filename(self, name: str) -> str:
        cleaned = re.sub(r"[^A-Za-z0-9._-]+", "_", (name or "").strip())
        cleaned = cleaned.strip("._-")
        return cleaned or "trace"

    def _build_relative_trace_path(self, trace_id: str) -> str:
        filename = f"{self._sanitize_filename(trace_id)}.json"
        return f"{self.storage_prefix}/{filename}"

    def save_trace(self, trace_id: str, payload: Dict) -> Optional[str]:
        raise NotImplementedError

    def list_traces(self) -> List[Dict]:
        raise NotImplementedError

    def delete_trace(self, trace_id: str) -> bool:
        raise NotImplementedError

    def delete_all(self) -> int:
        raise NotImplementedError


class WorkloadTracesFolderService(BaseTracesFolderService):
    """
    Workload-local folder implementation.
    """

    def __init__(self, logger, storage_prefix: str):
        super().__init__(logger=logger, storage_prefix=storage_prefix)
        self._workload_folder = self._resolve_workload_dir()
        self._workload_folder.mkdir(parents=True, exist_ok=True)
        self.logger.debug(
            "Using workload folder storage for pasted traces: %s",
            str(self._workload_folder),
        )

    def _resolve_workload_dir(self) -> Path:
        try:
            from dataiku.core.workload_local_folder import get_workload_local_folder_path

            root = Path(get_workload_local_folder_path())
            return root / self.storage_prefix
        except Exception:
            fallback = Path("/tmp/traces-explorer-workload") / self.storage_prefix
            self.logger.warn(
                "Falling back to local storage directory for pasted traces: %s",
                str(fallback),
            )
            return fallback

    def _build_filepath(self, trace_id: str) -> Path:
        filename = f"{self._sanitize_filename(trace_id)}.json"
        return self._workload_folder / filename

    def save_trace(self, trace_id: str, payload: Dict) -> Optional[str]:
        target = self._build_filepath(trace_id)
        target.write_text(json.dumps(payload), encoding="utf-8")
        self.logger.debug(
            "Created pasted trace '%s' in workload path '%s'",
            trace_id,
            str(target),
        )
        return None

    def list_traces(self) -> List[Dict]:
        self.logger.debug("Listing pasted traces in workload folder '%s'", str(self._workload_folder))
        traces: List[Dict] = []
        for path in sorted(self._workload_folder.glob("*.json")):
            try:
                payload = json.loads(path.read_text(encoding="utf-8"))
                payload["__file_path"] = str(path)
                traces.append(payload)
                self.logger.debug("Listed pasted trace at workload path '%s'", str(path))
            except Exception as exc:
                self.logger.warn(
                    "Skipping invalid pasted trace file '%s': %s", str(path), exc
                )
        return traces

    def delete_trace(self, trace_id: str) -> bool:
        direct_path = self._build_filepath(trace_id)
        if not direct_path.exists():
            return False

        direct_path.unlink(missing_ok=True)
        self.logger.debug(
            "Deleted pasted trace '%s' from workload path '%s'",
            trace_id,
            str(direct_path),
        )
        return True

    def delete_all(self) -> int:
        deleted_count = 0
        for path in self._workload_folder.glob("*.json"):
            path.unlink(missing_ok=True)
            deleted_count += 1
            self.logger.debug("Deleted pasted trace from workload path '%s'", str(path))
        return deleted_count


class ManagedTracesFolderService(BaseTracesFolderService):
    """
    Dataiku managed-folder implementation.
    """

    def __init__(self, logger, managed_folder, storage_prefix: str):
        super().__init__(logger=logger, storage_prefix=storage_prefix)
        self._managed_folder = managed_folder

    def _list_managed_paths(self) -> List[str]:
        prefix = f"{self.storage_prefix}/"
        try:
            paths = self._managed_folder.list_paths_in_partition()
        except Exception:
            self.logger.warn(
                "Failed to list paths in managed folder for prefix '%s'", prefix
            )
            return []

        filtered: List[str] = []
        for raw_path in paths:
            path = str(raw_path).lstrip("/")
            if path.startswith(prefix) and path.endswith(".json"):
                filtered.append(path)
        return filtered

    def _read_managed_json(self, path: str) -> Dict:
        with self._managed_folder.get_download_stream(path) as stream:
            return json.loads(stream.read().decode("utf-8"))

    def _write_managed_json(self, path: str, payload: Dict) -> None:
        with self._managed_folder.get_writer(path) as writer:
            writer.write(json.dumps(payload).encode("utf-8"))

    def _delete_managed_path(self, path: str) -> bool:
        try:
            self._managed_folder.delete_path(path)
            return True
        except Exception:
            return False

    def save_trace(self, trace_id: str, payload: Dict) -> Optional[str]:
        target = self._build_relative_trace_path(trace_id)
        self._write_managed_json(target, payload)
        self.logger.debug(
            "Created pasted trace '%s' in managed folder path '%s'",
            trace_id,
            target,
        )
        return None

    def list_traces(self) -> List[Dict]:
        self.logger.debug("Listing pasted traces in managed folder '%s'", self.storage_prefix)
        traces: List[Dict] = []
        for path in sorted(self._list_managed_paths()):
            try:
                payload = self._read_managed_json(path)
                payload["__file_path"] = path
                traces.append(payload)
                self.logger.debug(
                    "Listed pasted trace at managed folder path '%s'",
                    path,
                )
            except Exception as exc:
                self.logger.warn(
                    "Skipping invalid pasted trace file '%s': %s", path, exc
                )
        return traces

    def delete_trace(self, trace_id: str) -> bool:
        target = self._build_relative_trace_path(trace_id)
        existing = set(self._list_managed_paths())
        if target not in existing:
            return False

        deleted = self._delete_managed_path(target)
        if not deleted:
            self.logger.warn(
                "Unable to delete managed-folder trace '%s': no supported delete API",
                target,
            )
            return False

        self.logger.debug(
            "Deleted pasted trace '%s' from managed folder path '%s'",
            trace_id,
            target,
        )
        return True

    def delete_all(self) -> int:
        deleted_count = 0
        for path in self._list_managed_paths():
            if self._delete_managed_path(path):
                deleted_count += 1
                self.logger.debug(
                    "Deleted pasted trace from managed folder path '%s'",
                    path,
                )
            else:
                self.logger.warn(
                    "Unable to delete managed-folder trace '%s': no supported delete API",
                    path,
                )
        return deleted_count


class NoStorageTracesFolderService(BaseTracesFolderService):
    """
    In-memory no-storage implementation.
    Keeps only the last pasted/imported trace payload.
    """

    def __init__(self, logger, storage_prefix: str):
        super().__init__(logger=logger, storage_prefix=storage_prefix)
        self._last_trace_payload: Optional[Dict] = None
        self.logger.debug(
            "Using no-storage mode for pasted traces (in-memory last trace only)"
        )

    def save_trace(self, trace_id: str, payload: Dict) -> Optional[str]:
        previous_trace_id = None
        if isinstance(self._last_trace_payload, dict):
            previous_trace_id = self._last_trace_payload.get("id")

        self._last_trace_payload = payload
        self.logger.debug(
            "Stored last pasted trace '%s' in memory (no persistent storage)",
            trace_id,
        )
        return str(previous_trace_id) if previous_trace_id else None

    def list_traces(self) -> List[Dict]:
        if self._last_trace_payload is None:
            return []
        return [self._last_trace_payload]

    def delete_trace(self, trace_id: str) -> bool:
        if self._last_trace_payload is None:
            return False

        if str(self._last_trace_payload.get("id")) != str(trace_id):
            return False

        self._last_trace_payload = None
        self.logger.debug(
            "Deleted in-memory pasted trace '%s' in no-storage mode", trace_id
        )
        return True

    def delete_all(self) -> int:
        if self._last_trace_payload is None:
            return 0
        self._last_trace_payload = None
        self.logger.debug("Deleted in-memory pasted trace in no-storage mode")
        return 1
