"""
Direct Access Resources - Resources class for Direct Service Access.
This class overrides methods to use direct HTTP requests instead of Snowflake service functions.
"""
from __future__ import annotations
from typing import Any, Dict, List, Optional, Union
import requests

from .... import debugging
from ....tools.constants import USE_GRAPH_INDEX, DEFAULT_QUERY_TIMEOUT_MINS, Generation
from ....environments import runtime_env, SnowbookEnvironment
from ...config import Config, ConfigStore, ENDPOINT_FILE
from ...direct_access_client import DirectAccessClient
from ...types import EngineState
from ...util import get_pyrel_version, poll_with_specified_overhead, safe_json_loads, ms_to_timestamp
from ....errors import ResponseStatusException, QueryTimeoutExceededException
from snowflake.snowpark import Session

# Import UseIndexResources to enable use_index functionality with direct access
from .use_index_resources import UseIndexResources

# Import helper functions from util
from .util import is_engine_issue as _is_engine_issue, is_database_issue as _is_database_issue, collect_error_messages

from .use_index_poller import DirectUseIndexPoller
from typing import Iterable

# Constants
TXN_ABORT_REASON_TIMEOUT = "transaction timeout"


class DirectAccessResources(UseIndexResources):
    """
    Resources class for Direct Service Access avoiding Snowflake service functions.
    Uses HTTP requests instead of Snowflake SQL for execution.
    """
    def __init__(
        self,
        profile: Union[str, None] = None,
        config: Union[Config, None] = None,
        connection: Union[Session, None] = None,
        dry_run: bool = False,
        reset_session: bool = False,
        generation: Optional[Generation] = None,
        language: str = "rel",
    ):
        super().__init__(
            generation=generation,
            profile=profile,
            config=config,
            connection=connection,
            reset_session=reset_session,
            dry_run=dry_run,
            language=language,
        )
        self._endpoint_info = ConfigStore(ENDPOINT_FILE)
        self._service_endpoint = ""
        self._direct_access_client = None
        # database and language are already set by UseIndexResources.__init__

    @property
    def service_endpoint(self) -> str:
        return self._retrieve_service_endpoint()

    def _retrieve_service_endpoint(self, enforce_update=False) -> str:
        account = self.config.get("account")
        app_name = self.config.get("rai_app_name")
        service_endpoint_key = f"{account}.{app_name}.service_endpoint"
        if self._service_endpoint and not enforce_update:
            return self._service_endpoint
        if self._endpoint_info.get(service_endpoint_key, "") and not enforce_update:
            self._service_endpoint = str(self._endpoint_info.get(service_endpoint_key, ""))
            return self._service_endpoint

        is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
        query = f"CALL {self.get_app_name()}.app.service_endpoint({not is_snowflake_notebook});"
        result = self._exec(query)
        assert result, f"Could not retrieve service endpoint for {self.get_app_name()}"
        if is_snowflake_notebook:
            self._service_endpoint = f"http://{result[0]['SERVICE_ENDPOINT']}"
        else:
            self._service_endpoint = f"https://{result[0]['SERVICE_ENDPOINT']}"

        self._endpoint_info.set(service_endpoint_key, self._service_endpoint)
        # save the endpoint to `ENDPOINT_FILE` to avoid calling the endpoint with every
        # pyrel execution
        try:
            self._endpoint_info.save()
        except Exception:
            print("Failed to persist endpoints to file. This might slow down future executions.")

        return self._service_endpoint

    @property
    def direct_access_client(self) -> DirectAccessClient:
        if self._direct_access_client:
            return self._direct_access_client
        try:
            service_endpoint = self.service_endpoint
            self._direct_access_client = DirectAccessClient(
                self.config, self.token_handler, service_endpoint, self.generation,
            )
        except Exception as e:
            raise e
        return self._direct_access_client

    def request(
        self,
        endpoint: str,
        payload: Dict[str, Any] | None = None,
        headers: Dict[str, str] | None = None,
        path_params: Dict[str, str] | None = None,
        query_params: Dict[str, str] | None = None,
        skip_auto_create: bool = False,
        skip_engine_db_error_retry: bool = False,
    ) -> requests.Response:
        with debugging.span("direct_access_request"):
            def _send_request():
                return self.direct_access_client.request(
                    endpoint=endpoint,
                    payload=payload,
                    headers=headers,
                    path_params=path_params,
                    query_params=query_params,
                )
            try:
                response = _send_request()
                if response.status_code != 200:
                    # For 404 responses with skip_auto_create=True, return immediately to let caller handle it
                    # (e.g., get_engine needs to check 404 and return None for auto_create_engine)
                    # For skip_auto_create=False, continue to auto-creation logic below
                    if response.status_code == 404 and skip_auto_create:
                        return response

                    try:
                        message = response.json().get("message", "")
                    except requests.exceptions.JSONDecodeError:
                        # Can't parse JSON response. For skip_auto_create=True (e.g., get_engine),
                        # this should have been caught by the 404 check above, so this is an error.
                        # For skip_auto_create=False, we explicitly check status_code below,
                        # so we don't need to parse the message.
                        if skip_auto_create:
                            raise ResponseStatusException(
                                f"Failed to parse error response from endpoint {endpoint}.", response
                            )
                        message = ""  # Not used when we check status_code directly

                    # fix engine on engine error and retry
                    # Skip setting up GI if skip_auto_create is True to avoid recursion or skip_engine_db_error_retry is true to let _exec_async_v2 perform the retry with the correct headers.
                    if ((_is_engine_issue(message) and not skip_auto_create) or _is_database_issue(message)) and not skip_engine_db_error_retry:
                        engine_name = payload.get("caller_engine_name", "") if payload else ""
                        engine_name = engine_name or self.get_default_engine_name()
                        engine_size = self.config.get_default_engine_size()
                        # Use the mixin's _poll_use_index method
                        self._poll_use_index(
                            app_name=self.get_app_name(),
                            sources=self.sources,
                            model=self.database,
                            engine_name=engine_name,
                            engine_size=engine_size,
                            headers=headers,
                        )
                        response = _send_request()
            except requests.exceptions.ConnectionError as e:
                messages = collect_error_messages(e)
                if any("nameresolutionerror" in msg for msg in messages):
                    # when we can not resolve the service endpoint, we assume it is outdated
                    # hence, we try to retrieve it again and query again.
                    self.direct_access_client.service_endpoint = self._retrieve_service_endpoint(
                        enforce_update=True,
                    )
                    return _send_request()
                # raise in all other cases
                raise e
            return response

    def _txn_request_with_gi_retry(
        self,
        payload: Dict,
        headers: Dict[str, str],
        query_params: Dict,
        engine: Union[str, None],
    ):
        """Make request with graph index retry logic.

        Attempts request with gi_setup_skipped=True first. If an engine or database
        issue occurs, polls use_index and retries with gi_setup_skipped=False.
        """
        response = self.request(
            "create_txn", payload=payload, headers=headers, query_params=query_params, skip_auto_create=True, skip_engine_db_error_retry=True
        )

        if response.status_code != 200:
            try:
                message = response.json().get("message", "")
            except requests.exceptions.JSONDecodeError:
                message = ""

            if _is_engine_issue(message) or _is_database_issue(message):
                engine_name = engine or self.get_default_engine_name()
                engine_size = self.config.get_default_engine_size()
                # Use the mixin's _poll_use_index method
                self._poll_use_index(
                    app_name=self.get_app_name(),
                    sources=self.sources,
                    model=self.database,
                    engine_name=engine_name,
                    engine_size=engine_size,
                    headers=headers,
                )
                headers['gi_setup_skipped'] = 'False'
                response = self.request(
                    "create_txn", payload=payload, headers=headers, query_params=query_params, skip_auto_create=True, skip_engine_db_error_retry=True
                )
            else:
                raise ResponseStatusException("Failed to create transaction.", response)

        return response

    def _exec_async_v2(
        self,
        database: str,
        engine: Union[str, None],
        raw_code: str,
        inputs: Dict | None = None,
        readonly=True,
        nowait_durable=False,
        headers: Dict[str, str] | None = None,
        bypass_index=False,
        language: str = "rel",
        query_timeout_mins: int | None = None,
        gi_setup_skipped: bool = False,
    ):

        with debugging.span("transaction") as txn_span:
            with debugging.span("create_v2") as create_span:

                use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)

                payload = {
                    "dbname": database,
                    "engine_name": engine,
                    "query": raw_code,
                    "v1_inputs": inputs,
                    "nowait_durable": nowait_durable,
                    "readonly": readonly,
                    "language": language,
                }
                if query_timeout_mins is None and (timeout_value := self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
                    query_timeout_mins = int(timeout_value)
                if query_timeout_mins is not None:
                    payload["timeout_mins"] = query_timeout_mins
                query_params={"use_graph_index": str(use_graph_index and not bypass_index)}

                # Add gi_setup_skipped to headers
                if headers is None:
                    headers = {}
                headers["gi_setup_skipped"] = str(gi_setup_skipped)
                headers['pyrel_program_id'] = debugging.get_program_span_id() or ""

                response = self._txn_request_with_gi_retry(
                    payload, headers, query_params, engine
                )

                artifact_info = {}
                response_content = response.json()

                txn_id = response_content["transaction"]['id']
                state = response_content["transaction"]['state']

                txn_span["txn_id"] = txn_id
                create_span["txn_id"] = txn_id
                debugging.event("transaction_created", txn_span, txn_id=txn_id)

            # fast path: transaction already finished
            if state in ["COMPLETED", "ABORTED"]:
                if txn_id in self._pending_transactions:
                    self._pending_transactions.remove(txn_id)

                # Process rows to get the rest of the artifacts
                for result in response_content.get("results", []):
                    filename = result['filename']
                    # making keys uppercase to match the old behavior
                    artifact_info[filename] = {k.upper(): v for k, v in result.items()}

            # Slow path: transaction not done yet; start polling
            else:
                self._pending_transactions.append(txn_id)
                with debugging.span("wait", txn_id=txn_id):
                    poll_with_specified_overhead(
                        lambda: self._check_exec_async_status(txn_id, headers=headers), 0.1
                    )
                artifact_info = self._list_exec_async_artifacts(txn_id, headers=headers)

            with debugging.span("fetch"):
                return self._download_results(artifact_info, txn_id, state)

    def _prepare_index(
        self,
        model: str,
        engine_name: str,
        engine_size: str = "",
        language: str = "rel",
        rai_relations: List[str] | None = None,
        pyrel_program_id: str | None  = None,
        skip_pull_relations: bool = False,
        headers: Dict | None = None,
    ):
        """
        Prepare the index for the given engine and model.
        """
        with debugging.span("prepare_index"):
            if headers is None:
                headers = {}

            payload = {
                "model_name": model,
                "caller_engine_name": engine_name,
                "language": language,
                "pyrel_program_id": pyrel_program_id,
                "skip_pull_relations": skip_pull_relations,
                "rai_relations": rai_relations or [],
                "user_agent": get_pyrel_version(self.generation),
            }
            # Only include engine_size if it has a non-empty string value
            if engine_size and engine_size.strip():
                payload["caller_engine_size"] = engine_size

            response = self.request(
                "prepare_index", payload=payload, headers=headers
            )

            if response.status_code != 200:
                raise ResponseStatusException("Failed to prepare index.", response)

            return response.json()

    def _check_exec_async_status(self, txn_id: str, headers: Dict[str, str] | None = None) -> bool:
        """Check whether the given transaction has completed."""

        with debugging.span("check_status"):
            response = self.request(
                "get_txn",
                headers=headers,
                path_params={"txn_id": txn_id},
            )
            assert response, f"No results from get_transaction('{txn_id}')"

        response_content = response.json()
        transaction = response_content["transaction"]
        status: str = transaction['state']

        # remove the transaction from the pending list if it's completed or aborted
        if status in ["COMPLETED", "ABORTED"]:
            if txn_id in self._pending_transactions:
                self._pending_transactions.remove(txn_id)

        if status == "ABORTED" and transaction.get("abort_reason", "") == TXN_ABORT_REASON_TIMEOUT:
            config_file_path = getattr(self.config, 'file_path', None)
            timeout_ms = int(transaction.get("timeout_ms", 0))
            timeout_mins = timeout_ms // 60000 if timeout_ms > 0 else int(self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS) or DEFAULT_QUERY_TIMEOUT_MINS)
            raise QueryTimeoutExceededException(
                timeout_mins=timeout_mins,
                query_id=txn_id,
                config_file_path=config_file_path,
            )

        # @TODO: Find some way to tunnel the ABORT_REASON out. Azure doesn't have this, but it's handy
        return status == "COMPLETED" or status == "ABORTED"

    def _list_exec_async_artifacts(self, txn_id: str, headers: Dict[str, str] | None = None) -> Dict[str, Dict]:
        """Grab the list of artifacts produced in the transaction and the URLs to retrieve their contents."""
        with debugging.span("list_results"):
            response = self.request(
                "get_txn_artifacts",
                headers=headers,
                path_params={"txn_id": txn_id},
            )
            assert response, f"No results from get_transaction_artifacts('{txn_id}')"
            artifact_info = {}
            for result in response.json()["results"]:
                filename = result['filename']
                # making keys uppercase to match the old behavior
                artifact_info[filename] = {k.upper(): v for k, v in result.items()}
            return artifact_info

    def get_transaction_problems(self, txn_id: str) -> List[Dict[str, Any]]:
        with debugging.span("get_transaction_problems"):
            response = self.request(
                "get_txn_problems",
                path_params={"txn_id": txn_id},
            )
            response_content = response.json()
            if not response_content:
                return []
            return response_content.get("problems", [])

    def get_transaction_events(self, transaction_id: str, continuation_token: str = ''):
        response = self.request(
            "get_txn_events",
            path_params={"txn_id": transaction_id, "stream_name": "profiler"},
            query_params={"continuation_token": continuation_token},
        )
        response_content = response.json()
        if not response_content:
            return {
                "events": [],
                "continuation_token": None
            }
        return response_content

    #--------------------------------------------------
    # Databases
    #--------------------------------------------------

    def get_installed_packages(self, database: str) -> Union[Dict, None]:
        use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
        if use_graph_index:
            response = self.request(
                "get_model_package_versions",
                payload={"model_name": database},
            )
        else:
            response = self.request(
                "get_package_versions",
                path_params={"db_name": database},
            )
        if response.status_code == 404 and response.json().get("message", "") == "database not found":
            return None
        if response.status_code != 200:
            raise ResponseStatusException(
                f"Failed to retrieve package versions for {database}.", response
            )

        content = response.json()
        if not content:
            return None

        return safe_json_loads(content["package_versions"])

    def get_database(self, database: str):
        with debugging.span("get_database", dbname=database):
            if not database:
                raise ValueError("Database name must be provided to get database.")
            response = self.request(
                "get_db",
                path_params={},
                query_params={"name": database},
            )
            if response.status_code != 200:
                raise ResponseStatusException(f"Failed to get db. db:{database}", response)

            response_content = response.json()

            if (response_content.get("databases") and len(response_content["databases"]) == 1):
                db = response_content["databases"][0]
                return {
                    "id": db["id"],
                    "name": db["name"],
                    "created_by": db.get("created_by"),
                    "created_on": ms_to_timestamp(db.get("created_on")),
                    "deleted_by": db.get("deleted_by"),
                    "deleted_on": ms_to_timestamp(db.get("deleted_on")),
                    "state": db["state"],
                }
            else:
                return None

    def create_graph(self, name: str):
        with debugging.span("create_model", dbname=name):
            return self._create_database(name,"")

    def delete_graph(self, name:str, force=False, language: str = "rel"):
        prop_hdrs = debugging.gen_current_propagation_headers()
        if self.config.get("use_graph_index", USE_GRAPH_INDEX):
            keep_database = not force and self.config.get("reuse_model", True)
            with debugging.span("release_index", name=name, keep_database=keep_database, language=language):
                response = self.request(
                    "release_index",
                    payload={
                        "model_name": name,
                        "keep_database": keep_database,
                        "language": language,
                        "user_agent": get_pyrel_version(self.generation),
                    },
                    headers=prop_hdrs,
                )
                if (
                    response.status_code != 200
                    and not (
                        response.status_code == 404
                        and "database not found" in response.json().get("message", "")
                    )
                ):
                    raise ResponseStatusException(f"Failed to release index. Model: {name} ", response)
        else:
            with debugging.span("delete_model", name=name):
                self._delete_database(name, headers=prop_hdrs)

    def clone_graph(self, target_name:str, source_name:str, nowait_durable=True, force=False):
        if force and self.get_graph(target_name):
            self.delete_graph(target_name)
        with debugging.span("clone_model", target_name=target_name, source_name=source_name):
            return self._create_database(target_name,source_name)

    def _delete_database(self, name:str, headers:Dict={}):
        with debugging.span("_delete_database", dbname=name):
            response = self.request(
                "delete_db",
                path_params={"db_name": name},
                query_params={},
                headers=headers,
            )
            if response.status_code != 200:
                raise ResponseStatusException(f"Failed to delete db. db:{name} ", response)

    def _create_database(self, name:str, source_name:str):
        with debugging.span("_create_database", dbname=name):
            payload = {
                "name": name,
                "source_name": source_name,
            }
            response = self.request(
                "create_db", payload=payload, headers={}, query_params={},
            )
            if response.status_code != 200:
                raise ResponseStatusException(f"Failed to create db. db:{name}", response)

    #--------------------------------------------------
    # Engines
    #--------------------------------------------------

    def list_engines(self, state: str | None = None):
        response = self.request("list_engines")
        if response.status_code != 200:
            raise ResponseStatusException(
                "Failed to retrieve engines.", response
            )
        response_content = response.json()
        if not response_content:
            return []
        engines = [
            {
                "name": engine["name"],
                "id": engine["id"],
                "size": engine["size"],
                "state": engine["status"], # callers are expecting 'state'
                "created_by": engine["created_by"],
                "created_on": engine["created_on"],
                "updated_on": engine["updated_on"],
            }
            for engine in response_content.get("engines", [])
            if state is None or engine.get("status") == state
        ]
        return sorted(engines, key=lambda x: x["name"])

    def get_engine(self, name: str):
        response = self.request("get_engine", path_params={"engine_name": name, "engine_type": "logic"}, skip_auto_create=True)
        if response.status_code == 404: # engine not found return 404
            return None
        elif response.status_code != 200:
            raise ResponseStatusException(
                f"Failed to retrieve engine {name}.", response
            )
        engine = response.json()
        if not engine:
            return None
        engine_state: EngineState = {
            "name": engine["name"],
            "id": engine["id"],
            "size": engine["size"],
            "state": engine["status"], # callers are expecting 'state'
            "created_by": engine["created_by"],
            "created_on": engine["created_on"],
            "updated_on": engine["updated_on"],
            "version": engine["version"],
            "auto_suspend": engine["auto_suspend_mins"],
            "suspends_at": engine["suspends_at"],
        }
        return engine_state

    def _create_engine(
            self,
            name: str,
            size: str | None = None,
            auto_suspend_mins: int | None = None,
            is_async: bool = False,
            headers: Dict[str, str] | None = None
        ):
        # only async engine creation supported via direct access
        if not is_async:
            return super()._create_engine(name, size, auto_suspend_mins, is_async, headers=headers)
        payload:Dict[str, Any] = {
            "name": name,
        }
        if auto_suspend_mins is not None:
            payload["auto_suspend_mins"] = auto_suspend_mins
        if size is not None:
            payload["size"] = size
        response = self.request(
            "create_engine",
            payload=payload,
            path_params={"engine_type": "logic"},
            headers=headers,
            skip_auto_create=True,
        )
        if response.status_code != 200:
            raise ResponseStatusException(
                f"Failed to create engine {name} with size {size}.", response
            )

    def delete_engine(self, name:str, force:bool = False, headers={}):
        response = self.request(
            "delete_engine",
            path_params={"engine_name": name, "engine_type": "logic"},
            headers=headers,
            skip_auto_create=True,
        )
        if response.status_code != 200:
            raise ResponseStatusException(
                f"Failed to delete engine {name}.", response
            )

    def suspend_engine(self, name:str):
        response = self.request(
            "suspend_engine",
            path_params={"engine_name": name, "engine_type": "logic"},
            skip_auto_create=True,
        )
        if response.status_code != 200:
            raise ResponseStatusException(
                f"Failed to suspend engine {name}.", response
            )

    def resume_engine_async(self, name:str, headers={}):
        response = self.request(
            "resume_engine",
            path_params={"engine_name": name, "engine_type": "logic"},
            headers=headers,
            skip_auto_create=True,
        )
        if response.status_code != 200:
            raise ResponseStatusException(
                f"Failed to resume engine {name}.", response
            )
        return {}

    def _poll_use_index(
        self,
        app_name: str,
        sources: Iterable[str],
        model: str,
        engine_name: str,
        engine_size: str | None = None,
        program_span_id: str | None = None,
        headers: Dict | None = None,
    ):
        """Poll use_index to prepare indices for the given sources using DirectUseIndexPoller."""
        return DirectUseIndexPoller(
            self,
            app_name=app_name,
            sources=sources,
            model=model,
            engine_name=engine_name,
            engine_size=engine_size,
            language=self.language,
            program_span_id=program_span_id,
            headers=headers,
            generation=self.generation,
        ).poll()

    def maybe_poll_use_index(
        self,
        app_name: str,
        sources: Iterable[str],
        model: str,
        engine_name: str,
        engine_size: str | None = None,
        program_span_id: str | None = None,
        headers: Dict | None = None,
    ):
        """Only call poll() if there are sources to process and cache is not valid."""
        sources_list = list(sources)
        self.database = model
        if sources_list:
            poller = DirectUseIndexPoller(
                self,
                app_name=app_name,
                sources=sources_list,
                model=model,
                engine_name=engine_name,
                engine_size=engine_size,
                language=self.language,
                program_span_id=program_span_id,
                headers=headers,
                generation=self.generation,
            )
            # If cache is valid (data freshness has not expired), skip polling
            if poller.cache.is_valid():
                cached_sources = len(poller.cache.sources)
                total_sources = len(sources_list)
                cached_timestamp = poller.cache._metadata.get("cachedIndices", {}).get(poller.cache.key, {}).get("last_use_index_update_on", "")

                message = f"Using cached data for {cached_sources}/{total_sources} data streams"
                if cached_timestamp:
                    print(f"\n{message} (cached at {cached_timestamp})\n")
                else:
                    print(f"\n{message}\n")
            else:
                return poller.poll()

