"""Token-limit fallback for context extraction."""

from __future__ import annotations

import ast
import hashlib
import logging
from typing import TYPE_CHECKING

from .imports import add_needed_imports_from_module
from .models import (
    AllContextResults,
    CodeContextType,
    CodeOptimizationContext,
    CodeString,
    CodeStringsMarkdown,
    FileContextCache,
)
from .pruning import parse_code_and_prune_cst

if TYPE_CHECKING:
    from pathlib import Path

log = logging.getLogger(__name__)


def encoded_tokens_len(s: str) -> int:
    """Return the approximate token count for a string."""
    return int(len(s) * 0.25)


def re_extract_from_cache(
    file_caches: list[FileContextCache],
    context_type: CodeContextType,
    project_root: Path,
    *,
    remove_docstrings: bool = True,
) -> CodeStringsMarkdown:
    """Re-extract context from cached modules without file I/O."""
    strings: list[CodeString] = []
    for file_cache in file_caches:
        try:
            pruned = parse_code_and_prune_cst(
                file_cache.cleaned_module,
                context_type,
                file_cache.fto_names,
                file_cache.hoh_names,
                remove_docstrings=remove_docstrings,
            )
        except ValueError:
            continue
        if pruned.code.strip():
            if context_type == CodeContextType.HASHING:
                code = ast.unparse(ast.parse(pruned.code))
            else:
                code = add_needed_imports_from_module(
                    src_module_code=file_cache.original_module,
                    dst_module_code=pruned,
                    src_path=file_cache.file_path,
                    dst_path=file_cache.file_path,
                    project_root=project_root,
                    helper_fqns=file_cache.helper_fqns,
                    gathered_imports=file_cache.gathered_imports,
                )
            strings.append(
                CodeString(
                    code=code,
                    file_path=file_cache.relative_path,
                ),
            )
    return CodeStringsMarkdown(code_strings=strings)


def apply_token_limits(
    all_results: AllContextResults,
    project_root: Path,
    optim_token_limit: int = 30000,
) -> CodeOptimizationContext:
    """Apply progressive degradation when context exceeds token limits.

    Raises *ValueError* if read-writable context alone exceeds the limit.
    """
    rw_markdown = all_results.read_writable.markdown
    rw_tokens = encoded_tokens_len(rw_markdown)
    if rw_tokens > optim_token_limit:
        msg = "Read-writable context exceeds token limit"
        raise ValueError(msg)

    read_only = all_results.read_only.markdown
    ro_tokens = encoded_tokens_len(read_only)

    if rw_tokens + ro_tokens > optim_token_limit:
        log.debug(
            "Code context has exceeded token limit, "
            "removing docstrings from read-only code",
        )
        read_only = re_extract_from_cache(
            all_results.file_caches,
            CodeContextType.READ_ONLY,
            project_root,
            remove_docstrings=True,
        ).markdown
        if rw_tokens + encoded_tokens_len(read_only) > optim_token_limit:
            log.debug(
                "Code context has exceeded token limit, "
                "removing read-only code",
            )
            read_only = ""

    hashing_markdown = all_results.hashing.markdown
    return CodeOptimizationContext(
        read_writable=rw_markdown,
        read_only=read_only,
        hashing=hashing_markdown,
        testgen=all_results.testgen.markdown,
        hashing_hash=hashlib.sha256(
            hashing_markdown.encode("utf-8"),
        ).hexdigest(),
    )
