#
#   Imandra Inc.
#
#   api_ref.py
#

import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Final, Self
from uuid import UUID

import structlog
import typer

logger = structlog.get_logger(__name__)


@dataclass
class IMLAPIReference:
    id: UUID
    module: str
    name: str
    signature: str
    doc: str | None = None
    pattern: str | None = None

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Self:
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


_iar_cache: list[IMLAPIReference] | None = None

DOC_DIR: Final[Path] = Path(__file__).parent.parent / 'data'
SELECTED_GUIDE_DOC: Final[list[str]] = [
    'guides/tactics.md',
    'guides/ordinal.md',
    'guides/termination_proofs.md',
    'guides/region-decomp/quick-start.md',
    'guides/region-decomp/writing-decomp.md',
    'guides/region-decomp/how-to-use-decomp-for-analysis.md',
    'guides/region-decomp/examples/',
    'guides/verification-goal/quick-start.md',
    'guides/verification-goal/examples/',
]


@dataclass
class MDDoc:
    path: Path  # path relative to DOC_DIR
    title: str
    content: str


md_docs: list[MDDoc] | None = None


def _load_guide_docs() -> list[MDDoc]:
    """
    Load documents in guide/ directory.
    """
    global md_docs

    def get_doc(doc_rel_path: Path) -> MDDoc:
        content = (DOC_DIR / doc_rel_path).read_text()
        return MDDoc(
            doc_rel_path,
            doc_rel_path.stem,
            content,
        )

    if md_docs is None:
        md_docs = []
        for doc_path in SELECTED_GUIDE_DOC:
            abs_path = DOC_DIR / doc_path
            if abs_path.is_file():
                logger.debug(f'Loading guide doc file: {abs_path}')
                rel_path = abs_path.relative_to(DOC_DIR)
                md_docs.append(get_doc(rel_path))
            else:
                for doc_file in abs_path.glob('**/*.md'):
                    rel_path = doc_file.relative_to(DOC_DIR)
                    md_docs.append(get_doc(rel_path))
    return md_docs


def _load_iar() -> list[IMLAPIReference]:
    """
    Load IML API Reference data lazily.
    """
    global _iar_cache
    if _iar_cache is None:
        iar_path = DOC_DIR / 'iml_api_reference_202510011126.json'
        iar_data = json.loads(iar_path.read_text())
        _iar_cache = [IMLAPIReference.from_dict(i) for i in iar_data]
    return _iar_cache


def dump_iml_overview_and_api_refs(dir_path: Path):
    """
    Writes the IML overview and API reference to a directory.

    Output will be like this:
        - iml_overview.md
        - iml_api_reference/
            - global.md
            - module1.md
            - module2.md
            - etc...
    """
    if not dir_path.exists():
        dir_path.mkdir(exist_ok=False, parents=True)

    if not dir_path.is_dir():
        typer.secho(
            f'Path is not a directory: {dir_path}', fg=typer.colors.RED, err=True
        )
        raise typer.Exit(1)

    from codelogician.doc.utils.prompts import (
        iml_caveats,
        iml_intro,
        lang_agnostic_meta_eg_overview,
        lang_agnostic_meta_egs_str,
    )

    # Write overview
    iml_overview = (
        iml_intro
        + iml_caveats
        + lang_agnostic_meta_eg_overview
        + lang_agnostic_meta_egs_str
    )
    overview_path = dir_path / 'iml_overview.md'
    overview_path.write_text(iml_overview)
    typer.secho(f'✓ Written overview to {overview_path}', fg=typer.colors.GREEN)

    # Write API reference
    iar = _load_iar()
    api_ref_dir = dir_path / 'iml_api_reference'
    api_ref_dir.mkdir(exist_ok=True)

    # Group by module
    modules = sorted({entry.module for entry in iar})
    for module in modules:
        module_entries = [entry for entry in iar if entry.module == module]
        module_name = module or 'global'

        is_global: bool = module_name == 'global'

        # Create markdown content for this module
        module_lines = [
            f'# Module `{module_name}`\n' if not is_global else '# Global\n',
            f'\n{len(module_entries)} entries\n\n',
        ]

        if is_global:
            module_lines.append(
                f'The following entries are qualified by the module `{module_name}`.'
                '\n\n'
            )

        for entry in module_entries:
            module_lines.append(f'`{entry.name}`\n')
            module_lines.append(f'- Signature: `{entry.signature}`\n')
            if entry.doc:
                module_lines.append(f'- Doc: {entry.doc}.\n')
            # if entry.pattern:
            #     module_lines.append(f"\n**Pattern:** `{entry.pattern}`\n")
            module_lines.append('\n')

        # Write module file
        module_filename = f'{module_name}.md' if module else 'global.md'
        module_path = api_ref_dir / module_filename
        module_path.write_text(''.join(module_lines))

    typer.secho(
        f'✓ Written API reference to {api_ref_dir} ({len(modules)} modules)',
        fg=typer.colors.GREEN,
    )


def dump_guide_docs(dir_path: Path):
    """
    Writes the guide documents to a directory.

    - contract: The total number of dumped markdown files will always be >= len(SELECTED_GUIDE_DOC).
    This contract must be respected and is verified by tests.
    - deps: global SELECTED_GUIDE_DOC
    - effect: writes guide documents to dir_path

    Args:
        dir_path: The directory where guide documents will be written

    Output structure example:
        guides/
            tactics.md
            ordinal.md
            termination_proofs.md
            region-decomp/
                quick-start.md
                writing-decomp.md
                examples/
                    _index.md
                    basics/_index.md
                    ...
    """
    if not dir_path.exists():
        dir_path.mkdir(exist_ok=False, parents=True)

    if not dir_path.is_dir():
        typer.secho(
            f'Path is not a directory: {dir_path}', fg=typer.colors.RED, err=True
        )
        raise typer.Exit(1)

    guide_docs = _load_guide_docs()
    for doc in guide_docs:
        out_path = dir_path / doc.path
        out_path.parent.mkdir(exist_ok=True, parents=True)
        out_path.write_text(doc.content)
        typer.secho(f'✓ Written {doc.title} to {out_path}', fg=typer.colors.GREEN)
