"""Data types for the context extraction pipeline."""

from __future__ import annotations

import enum
import re
import sys
from pathlib import Path
from typing import TYPE_CHECKING, Any

import attrs

if sys.version_info >= (3, 11):
    from typing import Self
else:
    from typing_extensions import Self

if TYPE_CHECKING:
    import libcst as cst

    from .._model import FunctionParent, FunctionSource


_MARKDOWN_PATTERN = re.compile(
    r"```(\w+)(?::([^\n]+))?\n(.*?)\n```",
    re.DOTALL,
)


class CodeContextType(str, enum.Enum):
    """The four types of code context for optimization."""

    READ_WRITABLE = "READ_WRITABLE"
    READ_ONLY = "READ_ONLY"
    TESTGEN = "TESTGEN"
    HASHING = "HASHING"


@attrs.frozen
class CodeString:
    """A code snippet with an optional source file path."""

    code: str
    file_path: Path | None = None


@attrs.frozen
class CodeStringsMarkdown:
    """A collection of code snippets for context assembly."""

    code_strings: list[CodeString] = attrs.Factory(list)

    @property
    def markdown(self) -> str:
        """Format as markdown code blocks with file path suffixes."""
        return "\n".join(
            f"```python"
            f"{':' + cs.file_path.as_posix() if cs.file_path else ''}"
            f"\n{cs.code.strip()}\n```"
            for cs in self.code_strings
        )

    def file_to_code(self) -> dict[str, str]:
        """Return a mapping from file path (as string) to code."""
        return {str(cs.file_path): cs.code for cs in self.code_strings}

    def to_dict(self) -> dict[str, object]:
        """Serialize to a plain dictionary."""
        return {
            "code_strings": [
                {
                    "code": cs.code,
                    "file_path": (str(cs.file_path) if cs.file_path else None),
                }
                for cs in self.code_strings
            ],
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> Self:
        """Restore from a serialized dictionary."""
        raw: list[dict[str, object]] = data.get(  # type: ignore[assignment]
            "code_strings",
            [],
        )
        code_strings = []
        for cs in raw:
            fp = cs.get("file_path")
            code_strings.append(
                CodeString(
                    code=str(cs["code"]),
                    file_path=Path(str(fp)) if fp else None,
                ),
            )
        return cls(code_strings=code_strings)

    @staticmethod
    def parse_markdown_code(
        markdown_code: str,
    ) -> CodeStringsMarkdown:
        """Parse a Markdown string into a CodeStringsMarkdown object."""
        matches = _MARKDOWN_PATTERN.findall(markdown_code)
        code_string_list = []
        for _language, file_path, code in matches:
            path = Path(file_path.strip()) if file_path else None
            code_string_list.append(
                CodeString(code=code, file_path=path),
            )
        return CodeStringsMarkdown(code_strings=code_string_list)


@attrs.define(eq=False)
class UsageInfo:
    """Information about a name and its usage."""

    name: str
    used_by_qualified_function: bool = False
    dependencies: set[str] = attrs.Factory(set)


@attrs.frozen
class PruneConfig:
    """Controls per-context-type CST pruning behavior."""

    defs_with_usages: dict[str, UsageInfo] | None = None
    helpers: set[str] | None = None
    remove_docstrings: bool = False
    include_target_in_output: bool = True
    exclude_init_from_targets: bool = False
    keep_class_init: bool = False
    include_dunder_methods: bool = False
    include_init_dunder: bool = False


@attrs.frozen
class FileContextCache:
    """Per-file CST cache for re-extraction without re-reading."""

    original_module: cst.Module
    cleaned_module: cst.Module
    fto_names: set[str]
    hoh_names: set[str]
    helper_fqns: set[str]
    file_path: Path
    relative_path: Path
    gathered_imports: Any = None


@attrs.frozen
class CodeOptimizationContext:
    """The four context views for function optimization."""

    read_writable: str = ""
    read_only: str = ""
    hashing: str = ""
    testgen: str = ""
    hashing_hash: str = ""
    read_writable_code: CodeStringsMarkdown = attrs.Factory(
        CodeStringsMarkdown,
    )
    testgen_context: CodeStringsMarkdown = attrs.Factory(
        CodeStringsMarkdown,
    )
    helper_functions: list[FunctionSource] = attrs.Factory(list)
    testgen_helper_fqns: list[str] = attrs.Factory(list)
    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]] = (
        attrs.Factory(set)
    )

    def to_dict(self) -> dict[str, object]:
        """Serialize to a plain dictionary."""
        return {
            "read_writable": self.read_writable,
            "read_only": self.read_only,
            "hashing": self.hashing,
            "testgen": self.testgen,
            "hashing_hash": self.hashing_hash,
            "read_writable_code": self.read_writable_code.to_dict(),
            "testgen_context": self.testgen_context.to_dict(),
            "helper_functions": [h.to_dict() for h in self.helper_functions],
            "testgen_helper_fqns": list(self.testgen_helper_fqns),
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> Self:
        """Restore from a serialized dictionary."""
        from .._model import FunctionSource  # noqa: PLC0415

        rw_raw: dict[str, object] = data.get(  # type: ignore[assignment]
            "read_writable_code",
            {},
        )
        tg_raw: dict[str, object] = data.get(  # type: ignore[assignment]
            "testgen_context",
            {},
        )
        helpers_raw: list[dict[str, object]] = data.get(  # type: ignore[assignment]
            "helper_functions",
            [],
        )
        fqns_raw: list[str] = data.get(  # type: ignore[assignment]
            "testgen_helper_fqns",
            [],
        )
        return cls(
            read_writable=str(data.get("read_writable", "")),
            read_only=str(data.get("read_only", "")),
            hashing=str(data.get("hashing", "")),
            testgen=str(data.get("testgen", "")),
            hashing_hash=str(data.get("hashing_hash", "")),
            read_writable_code=CodeStringsMarkdown.from_dict(
                rw_raw,
            ),
            testgen_context=CodeStringsMarkdown.from_dict(
                tg_raw,
            ),
            helper_functions=[
                FunctionSource.from_dict(h) for h in helpers_raw
            ],
            testgen_helper_fqns=list(fqns_raw),
        )


@attrs.frozen
class AllContextResults:
    """Per-file context with caches for token-limit fallback."""

    read_writable: CodeStringsMarkdown = attrs.Factory(CodeStringsMarkdown)
    read_only: CodeStringsMarkdown = attrs.Factory(CodeStringsMarkdown)
    hashing: CodeStringsMarkdown = attrs.Factory(CodeStringsMarkdown)
    testgen: CodeStringsMarkdown = attrs.Factory(CodeStringsMarkdown)
    file_caches: list[FileContextCache] = attrs.Factory(list)
