from __future__ import annotations

import argparse
import importlib
import importlib.util
import pprint
import sys
from pathlib import Path
from types import ModuleType
from typing import Iterable

from botscore import constants as botsconfig


def _load_module(reference: str) -> ModuleType:
    path = Path(reference)
    if path.exists():
        spec = importlib.util.spec_from_file_location(path.stem, path)
        if spec is None or spec.loader is None:
            raise ImportError(f'Could not load module from {reference}.')
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        return module
    return importlib.import_module(reference)


def _collect_segments_from_structure(structure: list[dict]) -> set[str]:
    segments: set[str] = set()

    def walk(level: Iterable[dict]) -> None:
        for node in level:
            segment = node.get(botsconfig.ID)
            if segment:
                segments.add(segment)
            children = node.get(botsconfig.LEVEL)
            if children:
                walk(children)

    walk(structure)
    return segments


def _segments_from_grammar_module(reference: str) -> set[str]:
    module = _load_module(reference)
    structure = getattr(module, 'structure', None)
    if not isinstance(structure, list):
        raise ValueError(f'Grammar module {reference!r} does not expose a structure list.')
    return _collect_segments_from_structure(structure)


def _recorddefs_from_source(reference: str) -> dict[str, list]:
    module = _load_module(reference)
    recorddefs = getattr(module, 'recorddefs', None)
    if not isinstance(recorddefs, dict):
        raise ValueError(f'Recorddefs source {reference!r} does not expose a recorddefs dict.')
    return recorddefs


def _normalize_segments(raw_segments: Iterable[str]) -> list[str]:
    segments = {segment.strip() for segment in raw_segments if segment.strip()}
    return sorted(segments)


def _render_module(
    *,
    source_reference: str,
    grammar_references: list[str],
    segments: list[str],
    extracted: dict[str, list],
) -> str:
    grammar_comment = (
        ', '.join(grammar_references)
        if grammar_references
        else 'manual segment selection'
    )
    rendered_recorddefs = pprint.pformat(extracted, width=100, sort_dicts=True)
    return (
        '"""\n'
        'Generated by bots_airflow.devtools.extract_recorddefs.\n'
        f'Source: {source_reference}\n'
        f'Grammar selection: {grammar_comment}\n'
        f'Segments: {", ".join(segments)}\n'
        '"""\n\n'
        f'recorddefs = {rendered_recorddefs}\n'
    )


def build_recorddefs_module(
    *,
    source_reference: str,
    grammar_references: list[str] | None = None,
    segments: list[str] | None = None,
) -> tuple[list[str], dict[str, list], str]:
    source_recorddefs = _recorddefs_from_source(source_reference)
    segment_names: set[str] = set(segments or [])

    for grammar_reference in grammar_references or []:
        segment_names.update(_segments_from_grammar_module(grammar_reference))

    normalized_segments = _normalize_segments(segment_names)
    if not normalized_segments:
        raise ValueError('No segments were selected. Provide --grammar and/or --segment.')

    missing = [segment for segment in normalized_segments if segment not in source_recorddefs]
    if missing:
        raise KeyError(f'Segments missing from source recorddefs: {missing!r}')

    extracted = {segment: source_recorddefs[segment] for segment in normalized_segments}
    rendered = _render_module(
        source_reference=source_reference,
        grammar_references=list(grammar_references or []),
        segments=normalized_segments,
        extracted=extracted,
    )
    return normalized_segments, extracted, rendered


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description='Extract a minimal recorddefs module from a larger source module.',
    )
    parser.add_argument(
        '--source-recorddefs',
        required=True,
        help='Module path or file path exposing recorddefs.',
    )
    parser.add_argument(
        '--grammar',
        action='append',
        default=[],
        help='Grammar module path or file path to inspect for structure segments.',
    )
    parser.add_argument(
        '--segment',
        action='append',
        default=[],
        help='Explicit segment id to include. Can be passed multiple times.',
    )
    parser.add_argument(
        '--output',
        required=True,
        help='Target Python module file to write.',
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    segments, _recorddefs, rendered = build_recorddefs_module(
        source_reference=args.source_recorddefs,
        grammar_references=args.grammar,
        segments=args.segment,
    )
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(rendered)
    print(f'Wrote {len(segments)} segments to {output_path}')
    return 0


if __name__ == '__main__':
    raise SystemExit(main(sys.argv[1:]))
