import dspy
import os
from typing import Any
from modaic import PrecompiledConfig, PrecompiledProgram

lm = dspy.LM("openai/gpt-5-mini", max_tokens=16000)
lm_mini = dspy.LM("openai/gpt-5-mini", max_tokens=16000)
dspy.configure(lm=lm)


class DocWriter(dspy.Signature):
    """
    Write documentation for the provided Rails application source code.
    """

    source_tree: dict[str, Any] = dspy.InputField()
    documentation: str = dspy.OutputField(
        description="Generated markdown documentation."
    )


class DocWriterConfig(PrecompiledConfig):
    max_iterations: int = 10
    lm: str = "openai/gpt-5-mini"
    sub_lm: str = "openai/gpt-5-mini"


class DocWriterProgram(PrecompiledProgram):
    config: DocWriterConfig

    def __init__(self, config: DocWriterConfig, **kwargs):
        super().__init__(config=config, **kwargs)
        self.doc_writer = dspy.RLM(
            DocWriter,
            max_iterations=config.max_iterations,
            sub_lm=config.sub_lm,
            verbose=True,
        )

    def forward(self, source_root: str) -> dspy.Prediction:
        source_tree = load_source_tree(source_root)
        print(f"Loaded source tree with {len(source_tree)} entries")
        return self.doc_writer(source_tree=source_tree)


def load_source_tree(root_dir: str) -> dict[str, Any]:
    """Recursively load the folder into a nested dict."""
    tree: dict[str, Any] = {}
    for entry in os.listdir(root_dir):
        path = os.path.join(root_dir, entry)
        if os.path.isdir(path):
            tree[entry] = load_source_tree(path)
        else:
            with open(path, "r", encoding="utf-8", errors="ignore") as f:
                tree[entry] = f.read()
    return tree


doc_writer = DocWriterProgram(DocWriterConfig())
SOURCE_ROOT = "."


def main():
    """
    print("Starting documentation generation...")
    result = doc_writer(source_root=SOURCE_ROOT)

    with open("generated_documentation.md", "w", encoding="utf-8") as f:
        print("Writing documentation to file...")
        f.write(result.documentation)
        print("Documentation written to generated_documentation.md")

    print("Pushing documentation to hub...")
    """
    doc_writer.push_to_hub("farouk1/ruby-rails-doc-generator")
    print("Documentation generated and pushed to hub!")


if __name__ == "__main__":
    main()