"""Preprocessing step to extract dataclass constructor signatures.

When the test context contains dataclass definitions, this module extracts
the constructor signatures to help the LLM understand how to properly
instantiate them, including inherited fields from parent dataclasses.
"""

from __future__ import annotations

from typing import NamedTuple

import libcst as cst
from libcst.metadata import MetadataWrapper

from aiservice.common.cst_utils import get_base_class_name, get_node_source_text, has_decorator, parse_module_to_cst
from aiservice.common.markdown_utils import extract_all_code_from_markdown


class FieldInfo(NamedTuple):
    """Information about a dataclass field."""

    name: str
    annotation: str
    has_default: bool
    default_repr: str | None
    source: str  # "inherited" or "own"


def extract_dataclass_fields(
    class_node: cst.ClassDef, source_lines: list[str], wrapper: MetadataWrapper
) -> list[FieldInfo]:
    """Extract fields from a dataclass definition.

    Returns list of FieldInfo with (field_name, type_annotation, has_default, default_value).
    """
    fields: list[FieldInfo] = []
    for stmt in class_node.body.body:
        if not isinstance(stmt, cst.SimpleStatementLine):
            continue
        for node in stmt.body:
            if not (isinstance(node, cst.AnnAssign) and isinstance(node.target, cst.Name)):
                continue

            field_name = node.target.value
            annotation = get_node_source_text(node.annotation, source_lines, wrapper)

            default_repr = None
            if node.value is not None:
                default_repr = get_node_source_text(node.value, source_lines, wrapper)
                # Truncate long defaults
                if len(default_repr) > 50:
                    default_repr = default_repr[:47] + "..."

            fields.append(FieldInfo(field_name, annotation, node.value is not None, default_repr, "own"))
    return fields


class DataclassCollector(cst.CSTVisitor):
    """Visitor to collect all dataclass definitions."""

    def __init__(self) -> None:
        self.dataclasses: dict[str, cst.ClassDef] = {}

    def visit_ClassDef(self, node: cst.ClassDef) -> bool:
        if has_decorator(node, "dataclass"):
            self.dataclasses[node.name.value] = node
        return True  # Continue visiting nested classes


def find_all_dataclasses(source_code: str) -> tuple[dict[str, cst.ClassDef], list[str], MetadataWrapper | None]:
    """Find all dataclass definitions in source code.

    Returns tuple of (dict mapping class name to class_node, source_lines, metadata_wrapper).
    """
    try:
        tree = parse_module_to_cst(source_code)
    except cst.ParserSyntaxError:
        return {}, [], None

    wrapper = MetadataWrapper(tree)
    collector = DataclassCollector()
    wrapper.visit(collector)
    source_lines = source_code.split("\n")
    return collector.dataclasses, source_lines, wrapper


def get_all_fields_with_inheritance(
    class_name: str, dataclasses: dict[str, cst.ClassDef], source_lines: list[str], wrapper: MetadataWrapper
) -> list[FieldInfo]:
    """Get all fields for a dataclass, including inherited ones.

    Fields are returned in order: inherited first, then own fields.
    """
    if class_name not in dataclasses:
        return []

    class_node = dataclasses[class_name]

    # First, collect inherited fields from parent dataclasses
    inherited_fields: list[FieldInfo] = []
    if class_node.bases:
        for base in class_node.bases:
            base_name = get_base_class_name(base)
            if base_name and base_name in dataclasses:
                parent_fields = get_all_fields_with_inheritance(base_name, dataclasses, source_lines, wrapper)
                # Mark these as inherited
                inherited_fields.extend(
                    FieldInfo(field.name, field.annotation, field.has_default, field.default_repr, "inherited")
                    for field in parent_fields
                )

    # Then collect own fields
    own_fields = extract_dataclass_fields(class_node, source_lines, wrapper)

    return inherited_fields + own_fields


def format_constructor_signature(class_name: str, fields: list[FieldInfo]) -> str:
    """Format a constructor signature for documentation."""
    if not fields:
        return f"{class_name}() - no fields"

    # Separate required and optional fields
    required = [f for f in fields if not f.has_default]
    optional = [f for f in fields if f.has_default]

    lines = [f"Constructor signature for {class_name}:"]

    if required:
        lines.append("  Required (positional) arguments:")
        for f in required:
            source_note = " (from parent class)" if f.source == "inherited" else ""
            lines.append(f"    - {f.name}: {f.annotation}{source_note}")

    if optional:
        lines.append("  Optional (keyword) arguments:")
        for f in optional:
            default_note = f" = {f.default_repr}" if f.default_repr else " = ..."
            source_note = " (from parent class)" if f.source == "inherited" else ""
            lines.append(f"    - {f.name}: {f.annotation}{default_note}{source_note}")

    return "\n".join(lines)


def get_dataclass_constructor_notes(test_context: str) -> list[str]:
    """Generate notes about dataclass constructor signatures.

    Analyzes the test context for dataclass definitions and generates
    explicit notes about their constructor signatures, including inherited fields.

    Args:
        test_context: The source code context to analyze for dataclass definitions.

    """
    # Extract actual code from markdown if present
    code_to_analyze = extract_all_code_from_markdown(test_context) if "```python" in test_context else test_context

    # Find all dataclasses in the context
    dataclasses, source_lines, wrapper = find_all_dataclasses(code_to_analyze)

    if not dataclasses or wrapper is None:
        return []

    notes = []
    for class_name in dataclasses:
        fields = get_all_fields_with_inheritance(class_name, dataclasses, source_lines, wrapper)
        if fields:
            signature = format_constructor_signature(class_name, fields)
            notes.append(signature)

    return notes
