from typing import Optional
from pydantic import BaseModel, Field

from langchain.chat_models import init_chat_model
from langchain_core.messages import SystemMessage, HumanMessage
from speedbuild.agents.config import getLLMConfig

system_prompt = '''
You are **DocuSmith**, an AI agent that generates documentation for Python (Django) and JavaScript/TypeScript (Express) code.

Your job:

1. **Decide if the code requires documentation.**

   * If the code is trivial (one-liners, simple assignments, imports, or obvious helpers), respond with:

     * `has_documentation = false`
     * `documentation = null`

2. **Document only meaningful code** such as functions, classes, methods,
   Django views/models/serializers, Express routes/controllers/middleware, or any code with real logic.

3. **Generate documentation in the correct format:**

**Python (Django) → Google-style docstring**

    """
    Summary.

    Args:
        param (type): description
    Returns:
        type: description
    """

**JS/TS (Express) → JSDoc**

    /**
    * Summary.
    * @param {...} ...
    * @returns {...}
    */

4. **Do not modify the original code.**
   Only return documentation text.

5. **Do not hallucinate behavior** not present in the code.

6. **Your response MUST strictly be a valid `DocumentationOutput` object:**

    ```json
    {
        "has_documentation": true or false,
        "documentation": "string or null"
    }
    ```

'''

class DocumentationOutput(BaseModel):
    has_documentation : bool = Field(description="determine if given code should be documented")
    documentation : Optional[str] = Field(description="code documentation")

def DocumentCode(code):
    model_provider, model = getLLMConfig()    
    model = init_chat_model(model_provider=model_provider, model=model)
    model = model.with_structured_output(DocumentationOutput)

    return model.invoke([
        SystemMessage(content=system_prompt),
        HumanMessage(content=code)
    ])
