"""Storage layer — pluggable backends for traces.

Use get_store() to get the configured backend:

    from bandito.storage import get_store

    store = get_store()
    store.write_trace(trace)
    store.list_projects()

Storage config resolution:
1. Explicit args (backend=, directory=)
2. Environment variables (BANDITO_STORAGE_BACKEND, BANDITO_S3_BUCKET, etc.)
3. Local config.toml
4. Default: local → ~/.bandito/traces/
"""

from __future__ import annotations

import os

from .base import StoreBase, TraceReader, TraceStats
from .local import LocalStore
from .sqlite import SQLiteStore


class StorageError(Exception):
    """Raised when storage is misconfigured or unavailable."""
    pass


def get_store(
    backend: str | None = None,
    directory: str | None = None,
    project: str | None = None,
) -> TraceReader:
    """Resolve and return the configured store.

    Resolution order:
    1. Explicit args
    2. Environment variables (BANDITO_STORAGE_BACKEND)
    3. Org storage config from backend (if BANDITO_API_KEY set)
    4. Local config.toml
    5. Default: local

    Args:
        backend: Storage backend ("local", "s3", "postgres").
        directory: Override directory for local backend.
        project: Project name (used for local dir resolution).

    Returns:
        A TraceReader instance (StoreBase for writable backends,
        TraceReader for read-only backends like Langfuse).
    """
    backend, config = _resolve_backend_with_config(backend)
    return _build_store(backend, directory, project, config)


def _resolve_backend_with_config(explicit_backend: str | None = None) -> tuple[str, dict | None]:
    """Resolve storage backend and config from env, org config, or local config.
    
    Returns (backend, config) tuple. config may be None.
    """
    if explicit_backend:
        return explicit_backend, None

    env = os.environ.get("BANDITO_STORAGE_BACKEND")
    if env:
        return env, None

    api_key = os.environ.get("BANDITO_API_KEY")
    if api_key:
        try:
            from ..api import BanditoClient

            client = BanditoClient()
            org_storage = client.get_org_storage()
            if org_storage and org_storage.get("storage_backend"):
                return org_storage["storage_backend"], org_storage.get("storage_config")
        except Exception:
            pass

    try:
        from ..config import load_config

        config = load_config()
        if hasattr(config, "storage") and config.storage and config.storage.backend:
            return config.storage.backend, None
    except (FileNotFoundError, OSError):
        pass
    return "local", None


def _build_store(
    backend: str,
    directory: str | None = None,
    project: str | None = None,
    config: dict | None = None,
) -> TraceReader:
    """Build a store from backend type + config."""
    if backend == "local":
        return _make_sqlite_store(config or {})

    if backend == "sqlite":
        return _make_sqlite_store(config or {})

    if backend == "s3":
        return _make_s3_store(config or {})

    if backend == "postgres":
        return _make_postgres_store(config or {})

    if backend == "langfuse":
        return _make_langfuse_store(config or {}, project=project)

    raise StorageError(
        f"Unknown storage backend: '{backend}'. Supported: local, sqlite, s3, postgres, langfuse"
    )


def _resolve_local_dir(project: str | None = None) -> str:
    """Resolve local traces directory from env > config > default."""
    from pathlib import Path

    env = os.environ.get("BANDITO_TRACES_DIR")
    if env:
        return env
    try:
        from ..config import load_config
        config = load_config()
        if hasattr(config, "storage") and config.storage and config.storage.directory:
            return config.storage.directory
        return config.traces_dir
    except (FileNotFoundError, OSError):
        return str(Path.home() / ".bandito" / "traces")


def _make_sqlite_store(config: dict) -> TraceReader:
    """Create a SQLiteStore from config."""
    db_path = config.get("db_path") or os.environ.get("BANDITO_SQLITE_PATH")
    return SQLiteStore(db_path=db_path)


def _make_s3_store(config: dict) -> TraceReader:
    """Create an S3Store from config + env vars for credentials."""
    try:
        from .s3 import S3Store
    except ImportError:
        raise StorageError(
            "S3 backend requires boto3. Install with: pip install bandito[s3]"
        )

    bucket = config.get("bucket") or os.environ.get("BANDITO_S3_BUCKET")
    if not bucket:
        raise StorageError(
            "S3 backend requires a bucket. Configure via bandito setup config."
        )

    return S3Store(
        bucket=bucket,
        prefix=config.get("prefix", "") or os.environ.get("BANDITO_S3_PREFIX", ""),
        endpoint_url=config.get("endpoint") or os.environ.get("BANDITO_S3_ENDPOINT"),
        aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
        aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
        region_name=config.get("region") or os.environ.get("AWS_DEFAULT_REGION"),
    )


def _make_postgres_store(config: dict) -> TraceReader:
    """Create a PostgresStore from config + env vars."""
    try:
        from .postgres import PostgresStore
    except ImportError:
        raise StorageError(
            "PostgreSQL backend requires psycopg. Install with: pip install bandito[postgres]"
        )

    conninfo = config.get("url") or os.environ.get("BANDITO_POSTGRES_URL")
    if not conninfo:
        raise StorageError(
            "PostgreSQL backend requires a connection URL. Configure via bandito setup config."
        )

    return PostgresStore(conninfo=conninfo)


def _make_langfuse_store(config: dict, project: str | None = None) -> TraceReader:
    """Create a LangfuseStore from config + env vars.

    Resolution: explicit config > config.toml [storage.langfuse] > env vars.
    """
    try:
        from .langfuse import LangfuseStore
    except ImportError:
        raise StorageError(
            "Langfuse backend requires the langfuse package. Install with: pip install bandito[langfuse]"
        )

    # Try config.toml langfuse settings
    lf_config = {}
    try:
        from ..config import load_config
        cfg = load_config()
        if cfg.storage and cfg.storage.langfuse:
            lf = cfg.storage.langfuse
            if lf.public_key:
                lf_config["public_key"] = lf.public_key
            if lf.secret_key:
                lf_config["secret_key"] = lf.secret_key
            if lf.base_url:
                lf_config["base_url"] = lf.base_url
    except (FileNotFoundError, OSError):
        pass

    # Merge: explicit config > config.toml > env vars
    merged = {**lf_config, **config}

    return LangfuseStore(
        public_key=merged.get("public_key") or os.environ.get("LANGFUSE_PUBLIC_KEY"),
        secret_key=merged.get("secret_key") or os.environ.get("LANGFUSE_SECRET_KEY"),
        base_url=merged.get("base_url") or os.environ.get("LANGFUSE_BASE_URL"),
        project=project,
    )


__all__ = [
    "TraceReader",
    "StoreBase",
    "LocalStore",
    "TraceStats",
    "StorageError",
    "get_store",
]
