"""
Exercise 3 – Structured Input and Structured Output
====================================================
AISE501 · Prompting in Coding · Spring Semester 2026

Learning goals
--------------
* Request machine-parseable output (JSON and YAML) from the LLM.
* Parse the JSON response in Python and use it programmatically.
* Build a second prompt dynamically from the parsed data.
* Understand why structured output is essential for LLM pipelines.

Tasks
-----
Part A  Ask the LLM to review analyze_me.py and return a JSON report (TODOs 1-4).
Part B  Parse the JSON response and print a summary table (TODOs 5-6).
Part C  Use the parsed data to build a follow-up prompt automatically (TODOs 7-8).
Part D  Repeat Part A but request YAML instead of JSON (TODO 9).

Estimated time: 40-50 minutes
"""

import json
from pathlib import Path

from server_utils import chat, chat_json, get_client, print_messages, print_separator

client = get_client()

code_to_review = Path("analyze_me.py").read_text()


# ── Part A: Structured Input → JSON Output ────────────────────────────────────
print_separator("Part A – Request JSON Output")

# TODO 1:  Write a system prompt that instructs the model to ALWAYS respond
#          with valid JSON and nothing else (no markdown fences, no explanation).

system_a = """\
<request>
  <persona>You are a master python tutor</persona>
  <style>You follow the PEP 8 style guide</style>
  <constraints>Only respond in a json format following the user provided schema</constraints>
</request>
"""

# TODO 2:  Write the user prompt.
#          Use XML tags for <persona>, <task>, and <code>.
#
#          In <task>, specify the exact JSON schema you expect:
#
schema = """{
  "summary": "<one sentence overview>",
  "bugs": [
    {
      "id": 1,
      "severity": "Critical|Medium|Style",
      "line": <int or null>,
      "function": "<function name>",
      "description": "<what is wrong>",
      "fix": "<one-sentence fix hint>"
    },
    ...
  ],
  "overall_quality": "Poor|Fair|Good|Excellent"
}"""
#
# Tip: paste the schema directly inside a <schema> tag in your prompt.

prompt_a = f"""\
TODO: Write your structured prompt here.
Include <persona>, <task>, <schema>, and <code> tags.

<persona>
    You are a Python engineer who is rigorous about correctness and follows PEP-8 and best practices.
</persona>

<task>
    Review the Python code and identify ALL bugs.
    Explain all the bugs you found the schema provided.
</task>

<schema>
{schema}
</schema>

<code language="python" filename="analyze_me.py">
{code_to_review}
</code>"""

messages_a = [
    # TODO 3: build the messages list (system + user)
    {"role": "system", "content": system_a},
    {"role": "user",   "content": prompt_a},
]

# TODO 4: call chat_json() and store the raw response string in raw_json_a.
#          chat_json() adds response_format={"type": "json_object"} so the
#          server guarantees the output is parseable by json.loads().
print_messages(messages_a)
raw_json_a = chat_json(client, messages_a)
print("Raw response:")
print(raw_json_a)


# ── Part B: Parse the JSON and Display a Summary ──────────────────────────────
print_separator("Part B – Parse JSON and Print Summary")

# TODO 5:  Parse raw_json_a with json.loads().
#          Handle the case where the model returned malformed JSON
#          (wrap in try/except and print a helpful error message).

report = json.loads(raw_json_a)

# TODO 6:  Print a formatted summary table like this:
#
#   Overall quality : Fair
#   Summary         : ...
#
#   ID | Severity | Line | Function              | Description
#   ---+----------+------+-----------------------+---------------------------
#    1 | Critical |   12 | calculate_statistics  | ZeroDivisionError on ...
#    2 | ...
#
# Hint: use f-strings and ljust() / rjust() for alignment.

print(f"Overall quality : {report['overall_quality']}")
print(f"Summary         : {report['summary']}\n")

bugs = report.get("bugs", [])
if bugs:
    headers = {
        "id": "ID",
        "severity": "Severity",
        "line": "Line",
        "function": "Function",
        "description": "Description",
    }

    # Compute column widths
    widths = {
        key: max(len(headers[key]), *(len(str(b[key])) for b in bugs))
        for key in headers
    }

    # Header row
    print(
        f"{headers['id'].ljust(widths['id'])} | "
        f"{headers['severity'].ljust(widths['severity'])} | "
        f"{headers['line'].ljust(widths['line'])} | "
        f"{headers['function'].ljust(widths['function'])} | "
        f"{headers['description']}"
    )

    # Separator row
    print(
        f"{'-' * widths['id']}-+-"
        f"{'-' * widths['severity']}-+-"
        f"{'-' * widths['line']}-+-"
        f"{'-' * widths['function']}-+-"
        f"{'-' * widths['description']}"
    )

    # Data rows
    for bug in bugs:
        print(
            f"{str(bug['id']).ljust(widths['id'])} | "
            f"{bug['severity'].ljust(widths['severity'])} | "
            f"{str(bug['line']).ljust(widths['line'])} | "
            f"{bug['function'].ljust(widths['function'])} | "
            f"{bug['description']}"
        )

# ── Part C: Use the Parsed Data to Build a Follow-Up Prompt ──────────────────
print_separator("Part C – Dynamic Follow-Up Prompt from Parsed Data")

# TODO 7:  Select all bugs with severity "Critical" from the parsed report.
#          Build a new user prompt that:
#          - Lists each critical bug by ID and description
#          - Asks the LLM to provide the corrected code for each one
#          - Requests the output as a JSON OBJECT (not a bare array, because
#            response_format=json_object requires an object at the top level):
#            {"fixes": [{"bug_id": 1, "fixed_code": "..."}, ...]}
#
# Tip: wrap the schema in a {"fixes": [...]} object so chat_json() works.

critical_bugs = [b for b in report["bugs"] if b["severity"] == "Critical"]

followup_prompt = """\
TODO: Build the follow-up prompt dynamically using the critical_bugs list.
      Loop over critical_bugs to embed each bug's description in the prompt.
"""

# TODO 8:  Continue the conversation (multi-turn) by appending the previous
#          response and the new prompt, then call chat_json() and parse the result.
#          Because the schema is {"fixes": [...]}, extract the list with ["fixes"].

# messages_c = messages_a + [
#     {"role": "assistant", "content": raw_json_a},
#     {"role": "user",      "content": followup_prompt},
# ]
# print_messages(messages_c)
# raw_json_c = chat_json(client, messages_c)
# fixes = json.loads(raw_json_c)["fixes"]
# for fix in fixes:
#     print(f"\n--- Fix for bug {fix['bug_id']} ---")
#     print(fix["fixed_code"])


# ── Part D: Request YAML Instead of JSON ─────────────────────────────────────
print_separator("Part D – YAML Output")

# TODO 9:  Repeat Part A but ask for YAML output instead of JSON.
#          Install PyYAML if needed: pip install pyyaml
#          Parse the response with yaml.safe_load() and print the result.
#
# Question: Which format do you prefer for human-readable reports? For
#           machine-to-machine pipelines?

# import yaml
# ...


# ── Reflection Questions ──────────────────────────────────────────────────────
print_separator("Reflection Questions")
print(
    "1. What can go wrong when asking an LLM to return JSON?\n"
    "2. How did the <schema> tag influence the output structure?\n"
    "3. Why is structured output important for building LLM pipelines?\n"
    "4. When would you use JSON vs. YAML vs. plain text?\n"
)