from __future__ import annotations

from dataclasses import dataclass

import libcst as cst

from aiservice.common.cst_utils import any_ellipsis_in_cst, ellipsis_in_cst_not_types
from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context, split_markdown_code
from core.languages.python.testgen.preprocessing.preprocess_pipeline import preprocessing_testgen_pipeline


def format_notes_markdown(notes: list[str] | set[str]) -> str:
    """Format a collection of notes as markdown."""
    return f"Notes:\n{'\n'.join(notes)}" if notes else ""


@dataclass
class TestGenContextData:
    source_code_being_tested: str
    qualified_name: str


class BaseTestGenContext:
    def __init__(self, ctx_data: TestGenContextData) -> None:
        self.data = ctx_data

    @staticmethod
    def get_dynamic_context(ctx_data: TestGenContextData | None) -> BaseTestGenContext:
        if ctx_data is None:
            msg = "ctx_data cannot be None"
            raise ValueError(msg)
        if is_multi_context(ctx_data.source_code_being_tested):
            return MultiTestGenContext(ctx_data)
        return SingleTestGenContext(ctx_data)

    def validate_python_module(self, feature_version: tuple[int, int] | None = None) -> None:
        raise NotImplementedError

    def generate_notes_markdown(self) -> str:
        raise NotImplementedError

    def did_generate_ellipsis(self, generated_code: str, python_version: tuple[int, ...] | None = None) -> bool:
        raise NotImplementedError


class SingleTestGenContext(BaseTestGenContext):
    def validate_python_module(self, feature_version: tuple[int, int] | None = None) -> None:  # noqa: ARG002
        parsed = cst.parse_module(self.data.source_code_being_tested)
        if not parsed.body:
            msg = "Source code is empty or invalid"
            raise SyntaxError(msg)

    def generate_notes_markdown(self) -> str:
        return format_notes_markdown(preprocessing_testgen_pipeline(self.data.source_code_being_tested))

    def did_generate_ellipsis(self, generated_code: str, python_version: tuple[int, ...] | None = None) -> bool:  # noqa: ARG002
        original_module = cst.parse_module(self.data.source_code_being_tested)
        if any_ellipsis_in_cst(original_module):
            return False
        generated_module = cst.parse_module(generated_code)
        return ellipsis_in_cst_not_types(generated_module)


class MultiTestGenContext(BaseTestGenContext):
    def validate_python_module(self, feature_version: tuple[int, int] | None = None) -> None:  # noqa: ARG002
        for scoped_code in split_markdown_code(self.data.source_code_being_tested).values():
            parsed = cst.parse_module(scoped_code)
            if not parsed.body:
                msg = "Source code is empty or invalid"
                raise SyntaxError(msg)

    def generate_notes_markdown(self) -> str:
        all_notes: set[str] = set()
        for scoped_code in split_markdown_code(self.data.source_code_being_tested).values():
            all_notes.update(preprocessing_testgen_pipeline(scoped_code))
        return format_notes_markdown(all_notes)

    def did_generate_ellipsis(self, generated_code: str, python_version: tuple[int, ...] | None = None) -> bool:  # noqa: ARG002
        generated_module = cst.parse_module(generated_code)
        if not ellipsis_in_cst_not_types(generated_module):
            return False
        # Flag if any original section didn't have ellipsis (LLM introduced it)
        return any(
            not any_ellipsis_in_cst(cst.parse_module(scoped_code))
            for scoped_code in split_markdown_code(self.data.source_code_being_tested).values()
        )
