from modules import model, llm_g
import json
import re
from langchain_core.messages import SystemMessage, HumanMessage

def book_units_langchain_agent(
    topic: str,
    units_number: int,
    notes: str | None = None,
    course_description: str | None = None,
    course_audience: str | None = None,
):

    # ---------- Step 1 : Generate Markdown Outline ----------
    llm = llm_g()

    notes_block = ""
    if notes:
        notes_block = f"""
Additional author notes (optional guidance):
{notes}

Use these notes to guide the course structure but do NOT copy them literally.
"""

    description_block = ""
    if course_description:
        description_block = f"""
Book description (guidance):
{course_description}

Use it to understand the theme and depth of the course.
Do NOT copy it literally in the output.
"""

    audience_block = ""
    if course_audience:
        audience_block = f"""
Target audience (guidance):
{course_audience}

Use it to adapt the level and focus of the units.
"""

    result = llm.call(
        f"""
Create a structured outline for a professional course about [{topic}].

The course should contain exactly {units_number} units.

For each unit include:
- Unit title
- Learning outcome
- Main topics

{description_block}

{audience_block}

{notes_block}

Return the outline in Markdown.
Focus on logical learning progression.
"""
    )

    # ---------- Step 2 : System Prompt for JSON extraction ----------
    system_prompt = f"""
Extract structured course data from the Markdown text.

Return JSON with EXACTLY this structure:

{{
  "course_name": "{topic}",
  "course_description": "string",
  "course_audience": "string",
  "learning_outcomes": [
    "string"
  ],
  "units": [
    {{
      "unit_name": "string",
      "outcome": "string",
      "topics": [
        "string"
      ]
    }}
  ]
}}

Rules:

- If course_description is missing, generate one
- If course_audience is missing, generate one
- Generate 4-6 topics per unit
- Units must equal {units_number}
- unit_name should look like: "الفصل الأول: ..."
- topics must be short educational titles
- Output ONLY valid JSON
- No markdown
- No explanations
"""

    # ---------- Step 3 : Call Model ----------
    llm_model = model()

    response = llm_model.invoke(
        [
            SystemMessage(content=system_prompt),
            HumanMessage(content=result),
        ]
    )

    raw_output = response.content.strip()

    # ---------- Step 4 : Clean JSON ----------
    try:
        structured_data = json.loads(raw_output)
    except json.JSONDecodeError:
        cleaned = re.sub(r"```json|```", "", raw_output).strip()
        structured_data = json.loads(cleaned)

    return structured_data