from typing import Any
from .context import Context
from .itm import ExtractModel
from .executor import Response


class Extractor:
    def __init__(self, context: Context):
        self.context = context

    def extract(self, configs: list[ExtractModel], response: Response):
        for config in configs:
            value = self._get_value(response, config)
            self.context.set_extracted(config.name, value)

    def _get_value(self, response: Response, config: ExtractModel) -> Any:
        if config.source == "status":
            return response.status
        if config.source == "header":
            return response.headers.get(config.json_path)
        return self._json_path(response.body, config.json_path)

    def _json_path(self, data: dict, path: str) -> Any:
        parts = path.lstrip("$.").split(".")
        current = data
        for part in parts:
            if isinstance(current, dict):
                current = current.get(part)
            else:
                return None
        return current