"""Execution context for crawler operations."""

from __future__ import annotations

import os
import random
import shutil
from copy import deepcopy
from datetime import datetime, timezone
from io import BufferedReader, BytesIO
from pathlib import Path
from tempfile import mkdtemp
from typing import IO, TYPE_CHECKING, Any, BinaryIO, ContextManager, overload

from anystore.interface.rate_limit import RateLimit
from anystore.interface.tags import Tags
from anystore.logging import get_logger
from anystore.logic.serialize import to_store
from anystore.store import Store
from anystore.util import ensure_uuid
from anystore.util import join_relpaths as make_key
from followthemoney import EntityProxy, model
from ftm_lakehouse import get_archive, get_entities
from ftm_lakehouse.repository import ArchiveRepository, EntityRepository
from structlog.stdlib import BoundLogger

from memorious.core import get_cache, get_tags
from memorious.env import Env
from memorious.logic.check import ContextCheck
from memorious.settings import Settings
from memorious.util import make_url_key

if TYPE_CHECKING:
    from memorious.logic.crawler import Crawler
    from memorious.logic.http import ContextHttp
    from memorious.model.stage import CrawlerStage

DEFAULT_ORIGIN = "memorious"
CACHE_ORIGIN = "memorious-cache"


class BaseContext:
    """Base context providing storage, caching, and HTTP functionality.

    This class provides the shared functionality for both full crawler contexts
    and lightweight standalone fetch contexts.
    """

    dataset: str
    run_id: str
    work_path: str
    incremental: bool
    params: dict[str, Any]
    log: BoundLogger
    http: ContextHttp
    check: ContextCheck
    settings: Settings
    tags: Tags
    archive: ArchiveRepository
    cache: Store
    _stealthy: bool

    def __init__(
        self,
        dataset: str,
        run_id: str | None = None,
        incremental: bool = False,
        params: dict[str, Any] | None = None,
        stealthy: bool = False,
        logger_name: str | None = None,
    ):
        """Initialize the base context.

        Args:
            dataset: Dataset/crawler name used for namespacing storage.
            run_id: Unique run identifier (default: auto-generated UUID).
            incremental: Skip already-processed items (default: False).
            params: Configuration parameters (default: empty dict).
            stealthy: Use random user agents (default: False).
            logger_name: Name for the logger (default: dataset name).
        """
        from memorious.logic.http import ContextHttp

        self.dataset = dataset
        self.run_id = run_id or ensure_uuid()
        self.work_path = mkdtemp()
        self.incremental = incremental
        self.params = params or {}
        self._stealthy = stealthy

        self.log = get_logger(
            logger_name or dataset,
            dataset=dataset,
            run_id=self.run_id,
        )

        self.settings = Settings()
        self.env = Env()
        self.archive = get_archive(dataset)
        self.tags = get_tags(dataset)
        self.cache = get_cache()

        self.http = ContextHttp(self)
        self.check = ContextCheck(self)

    @property
    def stealthy(self) -> bool:
        """Whether to use random user agents for HTTP requests."""
        return self._stealthy

    def get(self, name: str, default: Any = None) -> Any:
        """Get a configuration value and expand environment variables."""
        value = self.params.get(name, default)
        if isinstance(value, str):
            value = os.path.expandvars(value)
        return value

    @overload
    def make_key(self, __part1: str, *parts: str, prefix: str | None = ...) -> str: ...

    @overload
    def make_key(self, *, prefix: str | None = ...) -> None: ...

    def make_key(self, *parts: str, prefix: str | None = None) -> str | None:
        """Create a namespaced key with the dataset name prefix.

        Args:
            *parts: Key parts to join. If empty/None, returns None.
            prefix: Optional prefix added after dataset name but not part
                of the None check. Useful for categorizing keys (e.g. "inc", "emit").

        Returns:
            Namespaced key string, or None if parts are empty.
        """
        key = make_key(*parts)
        if key is None:
            return None
        if key.startswith(self.dataset):  # erf
            key = key[len(self.dataset) + 1 :]
        if prefix:
            key = make_key(prefix, key)
        return make_key(self.dataset, key)

    def set_tag(self, key: str, value: Any = None) -> None:
        if not key or not key.strip():
            self.log.warning("Ignoring empty tag key")
            return
        if value is None:
            value = datetime.now(timezone.utc)
        self.tags.put(self.make_key(key), value)

    def get_tag(self, key: str) -> Any:
        if not key or not key.strip():
            self.log.warning("Ignoring empty tag key")
            return None
        return self.tags.get(self.make_key(key))

    def check_tag(self, key: str) -> bool:
        if not key or not key.strip():
            self.log.warning("Ignoring empty tag key")
            return False
        return self.tags.exists(self.make_key(key))

    def skip_incremental(self, *criteria: str) -> bool:
        """Perform an incremental check on a set of criteria.

        This can be used to execute a part of a crawler only once per an
        interval (which is specified by the ``expire`` setting). If the
        operation has already been performed (and should thus be skipped),
        this will return ``True``. If the operation needs to be executed,
        the returned value will be ``False``.
        """
        if not self.incremental:
            return False

        key = self.make_key(*criteria, prefix="inc")
        if key is None:
            return False

        if self.check_tag(key):
            return True

        self.set_tag(key, "inc")
        return False

    def store_file(
        self,
        file_path: str | Path,
        origin: str | None = DEFAULT_ORIGIN,
        checksum: str | None = None,
    ) -> str:
        """Put a file into permanent storage so it can be visible to other stages."""
        file_info = self.archive.store(
            file_path, checksum=checksum, origin=origin or DEFAULT_ORIGIN
        )
        return file_info.checksum

    def store_data(self, data: Any, checksum: str | None = None) -> str:
        if isinstance(data, (BinaryIO, BytesIO, BufferedReader)):
            fh = data
        else:
            fh = BytesIO(to_store(data))
        return self.archive.write_blob(fh, checksum=checksum)

    def open(self, content_hash: str) -> ContextManager[IO[bytes]]:
        return self.archive.open(content_hash)

    def local_path(self, content_hash: str) -> ContextManager[Path]:
        return self.archive.local_path(content_hash)

    def enforce_rate_limit(self, rate_limit: RateLimit) -> None:
        """
        Enforce rate limit for a resource.

        Updates the rate limit counter and blocks if limit exceeded.
        """
        rate_limit.update()
        if not rate_limit.check():
            rate_limit.comply()

    def close(self) -> None:
        """Clean up resources."""
        self.http.close()
        shutil.rmtree(self.work_path, ignore_errors=True)


class MemoriousContext(BaseContext):
    """Full crawler context with pipeline operations.

    Provides state tracking and methods for crawler operation interactions,
    including emit/recurse for pipeline stage transitions.
    """

    crawler: Crawler
    stage: CrawlerStage
    state: dict[str, Any]
    continue_on_error: bool | None
    entities: EntityRepository

    def __init__(self, crawler: Crawler, stage: CrawlerStage, state: dict[str, Any]):
        # Initialize base with crawler-derived values
        super().__init__(
            dataset=crawler.name,
            run_id=state.get("run_id"),
            incremental=state.get("incremental") or False,
            params=stage.params,
            stealthy=crawler.stealthy,
            logger_name=f"{crawler.name}.{stage.name}",
        )

        # Update logger with stage info
        self.log = get_logger(
            "%s.%s" % (crawler.name, stage.name),
            dataset=crawler.name,
            stage=stage.name,
            run_id=self.run_id,
        )

        # Crawler-specific attributes
        self.crawler = crawler
        self.stage = stage
        self.state = state
        self.continue_on_error = state.get("continue_on_error")

        # Entity storage (crawler-only)
        self.entities = get_entities(crawler.name)

    def _make_emit_cache_key(self, data: dict[str, Any]) -> str | None:
        """Generate a cache key for incremental emit tracking.

        Uses custom `emit_cache_key` or `foreign_id` if available, otherwise
        url, otherwise None.
        """
        cache_key = data.get("emit_cache_key")
        if cache_key:
            return self.make_key(cache_key, prefix="emit")
        foreign_id = data.get("foreign_id")
        if foreign_id:
            return self.make_key(foreign_id, prefix="emit")
        url = data.get("url")
        if url:
            return self.make_key(make_url_key(url), prefix="emit")
        return None

    def emit(
        self,
        rule: str = "pass",
        stage: str | None = None,
        data: dict[str, Any] | None = None,
        delay: int | None = None,
        optional: bool = False,
    ) -> None:
        """
        Invoke the next stage via procrastinate task queue.

        Args:
            rule: Handler rule name to determine target stage
            stage: Explicit target stage (overrides rule lookup)
            data: Data to pass to the next stage
            delay: Delay in seconds (logged but not applied - use scheduled_at for delays)
            optional: If True, silently skip if no target stage found
        """
        data = data or {}

        # Resolve target stage
        if stage is None:
            stage = self.stage.handlers.get(rule)
        if optional and stage is None:
            return
        if stage is None or stage not in self.crawler.stages:
            self.log.info("No next stage", stage=stage, rule=rule)
            return

        # Incremental: skip if already processed (cache key exists)
        cache_key = self._make_emit_cache_key(data)
        if self.incremental:
            if cache_key and self.check_tag(cache_key):
                self.log.info("Skipping emit (incremental)", cache_key=cache_key)
                return
        # Store cache key in data for marking complete at store stage
        data["_emit_cache_key"] = cache_key

        # Debug sampling
        if self.settings.debug:
            sampling_rate = self.get("sampling_rate")
            if sampling_rate and random.random() > float(sampling_rate):
                self.log.info("Skipping emit due to sampling rate", rate=sampling_rate)
                return

        # Log delay warning (procrastinate supports scheduled_at, but we simplify)
        if delay and delay > 0:
            self.log.debug("Delay requested but not applied", delay=delay)

        # Make a copy of the data to avoid mutation when in-memory connector
        data = deepcopy(data)

        # Defer the job via procrastinate
        self.crawler.defer(
            stage=stage,
            data=data,
            run_id=self.run_id,
            incremental=(
                self.incremental
                if self.incremental is not None
                else self.settings.incremental
            ),
            continue_on_error=self.continue_on_error or False,
        )

    def mark_emit_complete(self, data: dict[str, Any]) -> None:
        """Mark an emit cache key as complete.

        Called by store operations after successful storage to enable
        incremental skipping on future runs.
        """
        cache_key = data.get("_emit_cache_key")
        if cache_key:
            self.set_tag(cache_key, datetime.now())
            self.log.debug("Marked emit complete", cache_key=cache_key)

    def recurse(
        self, data: dict[str, Any] | None = None, delay: int | None = None
    ) -> None:
        """Have a stage invoke itself with a modified set of arguments."""
        if data is None:
            data = {}
        return self.emit(stage=self.stage.name, data=data, delay=delay)

    def execute(self, data: dict[str, Any]) -> Any:
        """Execute the stage method with the given data."""
        try:
            self.log.info(
                "Executing stage",
                method=self.stage.method_name,
            )
            return self.stage.method(self, data)
        except Exception as exc:
            self.emit_exception(exc)
            if not self.continue_on_error:
                raise exc
        finally:
            shutil.rmtree(self.work_path)

    def emit_warning(self, message: str, **kwargs: Any) -> None:
        self.log.warning(message, **kwargs)

    def emit_exception(self, exc: Exception) -> None:
        self.log.exception(str(exc))

    def dump_state(self) -> dict[str, Any]:
        state = deepcopy(self.state)
        state["dataset"] = self.crawler.name
        state["run_id"] = self.run_id
        return state

    @classmethod
    def from_state(
        cls, state: dict[str, Any], stage: str, config_file: str
    ) -> "MemoriousContext":
        """Create a Context from serialized state.

        Args:
            state: Serialized state dict containing dataset name and run_id
            stage: Stage name to execute
            config_file: Path or URI to crawler config file
        """
        from memorious.logic.crawler import get_crawler

        crawler = get_crawler(config_file)
        stage_obj = crawler.get(stage)
        if stage_obj is None:
            raise RuntimeError("[%r] has no stage: %s" % (crawler, stage))
        return cls(crawler, stage_obj, state)

    def make_entity(self, schema: str) -> EntityProxy:
        """Instantiate a FollowTheMoney entity"""
        return model.make_entity(
            schema, key_prefix=self.crawler.config.prefix or self.dataset
        )

    def __repr__(self) -> str:
        return "<Context(%r, %r)>" % (self.crawler, self.stage)


class FetchContext(BaseContext):
    """Lightweight context for standalone HTTP fetching without a crawler.

    This context provides the HTTP client, archive storage, and caching
    functionality without requiring a full crawler configuration.

    Example:
        >>> with FetchContext(dataset="my-scraper") as ctx:
        ...     response = ctx.http.get("https://example.com")
        ...     print(response.json)
    """

    def __init__(
        self,
        dataset: str = "fetch",
        cache: bool = True,
        proxies: str | list[str] | None = None,
        timeout: int | None = None,
        user_agent: str | None = None,
        stealthy: bool = False,
        incremental: bool = True,
    ):
        """Initialize a standalone fetch context.

        Args:
            dataset: Dataset name for archive namespace (default: "fetch").
            cache: Enable HTTP caching with 304 responses (default: True).
            proxies: Proxy URL or list of proxy URLs for rotation.
            timeout: HTTP request timeout in seconds.
            user_agent: Custom User-Agent header.
            stealthy: Use random user agents (default: False).
            incremental: Enable incremental state tracking (default: True).
        """
        # Build params from arguments
        params: dict[str, Any] = {"cache": cache}
        if proxies:
            params["http_proxies"] = proxies
        if timeout:
            params["http_timeout"] = timeout
        if user_agent:
            params["user_agent"] = user_agent

        super().__init__(
            dataset=dataset,
            incremental=incremental,
            params=params,
            stealthy=stealthy,
        )

    def __enter__(self) -> "FetchContext":
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()


# Backwards compatibility alias
Context = MemoriousContext
