"""Python-specific language state with lazy singleton accessors.

Caches expensive analysis artifacts that are shared across
the optimization session: module ASTs, reference graphs,
validated code, and function-to-test mappings.

Follows the platform-libs pattern where ``AzureCloudState``
provides ``resource_group()``, ``storage_account()``, etc.
as lazy-initialized singletons.  Here, ``PythonState``
provides ``reference_graph()``, ``module_ast(path)``, etc.
"""

from __future__ import annotations

import ast
import logging
from typing import TYPE_CHECKING

import attrs

from codeflash_core import LanguageState

from ._configuration import PythonConfiguration

if TYPE_CHECKING:
    from pathlib import Path

    from .analysis._reference_graph import ReferenceGraph
    from .pipeline._module_prep import ValidCode
    from .test_discovery.models import FunctionCalledInTest

log = logging.getLogger(__name__)


@attrs.define
class PythonState(LanguageState[PythonConfiguration]):
    """Cached analysis state for a Python optimization session.

    Each accessor lazily builds and caches its artifact on first
    access.  Subsequent calls return the cached instance.

    Use :meth:`close` to release resources (especially the
    reference graph's SQLite connection) when the session ends.
    """

    _reference_graph: ReferenceGraph | None = None
    _module_asts: dict[Path, ast.Module] = attrs.Factory(dict)
    _validated_code: dict[Path, ValidCode] = attrs.Factory(dict)
    _function_to_tests: dict[str, set[FunctionCalledInTest]] = attrs.Factory(
        dict
    )

    def reference_graph(self) -> ReferenceGraph:
        """Return the Jedi-based reference graph, building it if needed."""
        if self._reference_graph is None:
            from .analysis._reference_graph import (  # noqa: PLC0415
                ReferenceGraph,
            )

            self._reference_graph = ReferenceGraph(
                self.cfg.project_root,
            )
            log.debug("Built reference graph for %s", self.cfg.project_root)
        return self._reference_graph

    def module_ast(self, path: Path) -> ast.Module | None:
        """Return the parsed AST for *path*, caching the result.

        Returns *None* if the file contains invalid syntax.
        """
        if path not in self._module_asts:
            source = path.read_text(encoding="utf-8")
            try:
                self._module_asts[path] = ast.parse(
                    source,
                    filename=str(path),
                )
            except SyntaxError:
                log.warning("Failed to parse %s", path)
                return None
        return self._module_asts[path]

    def invalidate_module(self, path: Path) -> None:
        """Remove cached AST and validated code for *path*.

        Call this after applying an optimization to force
        re-parsing on the next access.
        """
        self._module_asts.pop(path, None)
        self._validated_code.pop(path, None)

    def validated_code(self, path: Path) -> ValidCode | None:
        """Return cached validated code for *path*, or *None*."""
        return self._validated_code.get(path)

    def set_validated_code(
        self,
        path: Path,
        valid: ValidCode,
    ) -> None:
        """Cache validated code for *path*."""
        self._validated_code[path] = valid

    def function_to_tests(
        self,
    ) -> dict[str, set[FunctionCalledInTest]]:
        """Return the function-to-test mapping."""
        return self._function_to_tests

    def set_function_to_tests(
        self,
        mapping: dict[str, set[FunctionCalledInTest]],
    ) -> None:
        """Replace the function-to-test mapping."""
        self._function_to_tests = mapping

    def close(self) -> None:
        """Release resources held by cached artifacts."""
        if self._reference_graph is not None:
            # ReferenceGraph holds a SQLite connection.
            self._reference_graph = None
        self._module_asts.clear()
        self._validated_code.clear()
        self._function_to_tests.clear()
        super().close()
