"""Policy definition tab UI components."""

import os
import sys

import gradio as gr

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.helpers import load_policy_from_file, load_preset_policy, list_preset_policy_names
from utils.policy_rules import (
    new_rule,
    policy_text_to_categories,
    parse_policy_markdown,
    rule_has_content,
    rule_is_enabled,
    policy_limits_markdown,
    render_rules_to_text,
    rules_have_content,
)


def build_policy_tab(base_dir: str, demo: gr.Blocks | None = None) -> dict:
    """Build the policy definition tab UI."""
    with gr.Tab("📋 Policy"):
        current_policy_state = gr.State(value="")
        uploaded_policies_state = gr.State(value={})  # Store uploaded policies: {"Uploaded - filename": content}
        # Structured rules: list of {"name","flag","dont_flag","exception"}. When these
        # carry content they are the source of truth — rendered to plain text for the
        # prompted models, and passed as a native category schema to any model that
        # accepts one (currently PolicyLM). Same policy either way.
        structured_rules_state = gr.State(value=[new_rule()])
        # Re-render counter. The rule editor must NOT re-render on every field edit:
        # each textbox writes back into structured_rules_state on blur, so a render
        # bound to that state tears the boxes down and rebuilds them under the
        # user's cursor, which reads as the fields being frozen. Only adding or
        # removing a rule changes which components exist, so only those bump this.
        rules_version_state = gr.State(value=0)
        # Name of whatever is currently loaded ("HS-P1 · …", an upload, or "" for
        # hand-written). The Testing tab shows this so nobody runs content against a
        # policy they have forgotten they loaded.
        policy_label_state = gr.State(value="")


        # --- Load a policy FIRST ------------------------------------------------------------
        # The common first action is "load my room's policy", so this sits above the rule
        # editor and opens by default; the blank editor used to be the first thing seen.
        with gr.Accordion("📥 Load a policy (start here)", open=True):
            gr.Markdown(
                "Pick a preset below to see the rules it contains, then add them to your policy. "
                "You can switch any rule on or off afterwards, and add more than one policy to "
                "compare or combine them."
            )
            with gr.Row():
                with gr.Column():
                    preset_dropdown = gr.Dropdown(
                        label="Select a policy to add",
                        choices=list_preset_policy_names(),
                        value=None
                    )
                    preset_preview = gr.Markdown("*Select a policy above to preview its rules.*")
                    load_preset_btn = gr.Button("➕ Add these rules to my policy", variant="primary")
                    credit_markdown = gr.Markdown("**Credit:** The original example policies came from the [Zentropi website](https://zentropi.ai/labelers). The **Workshop** presets are from the ROOST x Musubi policy-following workshop (27 Oct 2026).")

                with gr.Column():
                    gr.Markdown("Upload a markdown file:")
                    upload_file = gr.File(label="Upload Markdown File", file_types=[".md"])

        # --- What is actually loaded --------------------------------------------------------
        # Adding a preset used to change only the rule list further down the page, which
        # on a laptop is below the fold: people pressed the button and saw nothing move.
        # This line sits directly under the loader and names what is live right now.
        def _policy_status(rules, label):
            """One line naming the policy currently under test."""
            kept = [r for r in (rules or []) if rule_has_content(r)]
            if not kept:
                return "⚪ **No policy loaded** — add a preset above, or write your own rules below."
            on = [r for r in kept if rule_is_enabled(r)]
            sources = []
            for r in kept:
                # Just the ID ("MH-P1"), matching the tag in each rule's header.
                s = (r.get("source") or "").split(" · ")[0].strip() or "yours"
                if s not in sources:
                    sources.append(s)
            off_note = "" if len(on) == len(kept) else f" ({len(kept) - len(on)} switched off)"
            return (
                f"🟢 **Policy loaded:** {', '.join(sources)} — "
                f"**{len(on)}** of {len(kept)} rule{'' if len(kept) == 1 else 's'} active{off_note}."
            )

        policy_status = gr.Markdown(_policy_status([new_rule()], ""))

        # --- Structured rule editor (model-agnostic; native where supported) -------------------------------
        with gr.Accordion("🧩 Your policy as rules (switch any rule on or off)", open=True):
            gr.Markdown(
                "Each rule has three parts, which is how most written moderation policies are "
                "already structured:\n\n"
                "- **✅ Flag this**: what counts as a violation\n"
                "- **🚫 Don't flag this**: lookalikes that must *not* be flagged, such as "
                "counter-speech or quoting something in order to report it\n"
                "- **⚠️ Exception**: carve-outs to the violation rule\n\n"
                "Every enabled rule is sent to every model: as a category schema for models that "
                "accept one (a score **per rule**, with exceptions kept separate from violations), "
                "and as plain text for the rest. Switch a rule **off** to test without it."
            )
            limits_display = gr.Markdown(policy_limits_markdown([new_rule()]))
            add_rule_btn = gr.Button("➕ Add rule", variant="secondary")

            def _make_field_updater(idx: int, field: str):
                """Write a single edited field back into the rules list at idx."""
                def _upd(value, rules):
                    rules = [dict(r) for r in (rules or [])]
                    if 0 <= idx < len(rules):
                        rules[idx][field] = value
                    return rules
                return _upd

            def _make_header_updater(idx: int, field: str):
                """Write a field that appears in the accordion header, and redraw.

                Only the name and the on/off toggle do this. Both are set by an action
                that has already moved focus, so the re-render costs nothing, and
                without it the header would keep showing a stale name.
                """
                def _upd(value, rules, version):
                    rules = [dict(r) for r in (rules or [])]
                    if 0 <= idx < len(rules):
                        rules[idx][field] = value
                    return rules, (version or 0) + 1
                return _upd

            def _make_remover(idx: int):
                """Drop the rule at idx, keeping at least one (blank) rule."""
                def _rm(rules, version):
                    rules = [dict(r) for r in (rules or [])]
                    if 0 <= idx < len(rules):
                        rules.pop(idx)
                    return (rules or [new_rule()]), version + 1
                return _rm

            _render_triggers = [rules_version_state.change]
            if demo is not None:
                _render_triggers.append(demo.load)

            @gr.render(inputs=structured_rules_state, triggers=_render_triggers)
            def render_rules(rules):
                rules = rules or [new_rule()]
                for i, rule in enumerate(rules):
                    # One accordion per rule. A parsed policy can run to six rules and
                    # twenty-four textboxes; stacked flat that is unusable, so filled
                    # rules collapse to a titled header and only a blank one (just
                    # added) opens. The header carries the rule's own name, which is
                    # why renaming re-renders.
                    enabled = rule_is_enabled(rule)
                    rule_name = (rule.get("name") or "").strip()
                    # Just the ID ("HS-P1"), not the whole preset name — repeated down
                    # a six-rule list the full name drowns out the rule's own title.
                    origin = (rule.get("source") or "").strip().split(" · ")[0].strip()
                    # Say where the rule came from, so an edited preset rule and one
                    # you wrote yourself are not indistinguishable in a long list.
                    header = f"{i + 1}. {rule_name or 'Untitled rule'}   ·   {origin or 'yours'}"
                    if not enabled:
                        header += "   (off)"
                    with gr.Accordion(header, open=(not rule_has_content(rule)) or len(rules) == 1):
                        with gr.Row():
                            # interactive=True is explicit on every field below. Gradio
                            # normally infers it from whether a component is wired as an
                            # event input, but that inference does not hold for components
                            # built inside gr.render: the listeners are attached after the
                            # component is constructed, so the fields ship as display-only
                            # and cannot be typed into.
                            enabled_cb = gr.Checkbox(
                                label="Use this rule",
                                value=enabled,
                                scale=0,
                                min_width=140,
                                info="Switch off to test without it. The rule is kept.",
                                interactive=True,
                            )
                            remove_btn = gr.Button("🗑 Remove", scale=0, min_width=100, variant="secondary")
                        name_tb = gr.Textbox(
                            label="Rule name",
                            value=rule.get("name", ""),
                            placeholder="e.g. Threats and coercion",
                            interactive=True,
                        )
                        flag_tb = gr.Textbox(
                            label="✅ Flag this (violation)",
                            value=rule.get("flag", ""),
                            lines=2,
                            placeholder="What counts as a violation of this rule…",
                            interactive=True,
                        )
                        dont_tb = gr.Textbox(
                            label="🚫 Don't flag this",
                            value=rule.get("dont_flag", ""),
                            lines=2,
                            placeholder="Lookalikes that must NOT be flagged, e.g. counter-speech, quoting to report…",
                            interactive=True,
                        )
                        exc_tb = gr.Textbox(
                            label="⚠️ Exception",
                            value=rule.get("exception", ""),
                            lines=2,
                            placeholder="Carve-outs to the violation rule…",
                            interactive=True,
                        )
                        for field, comp in (
                            ("flag", flag_tb),
                            ("dont_flag", dont_tb),
                            ("exception", exc_tb),
                        ):
                            comp.blur(
                                _make_field_updater(i, field),
                                inputs=[comp, structured_rules_state],
                                outputs=structured_rules_state,
                            )
                        name_tb.blur(
                            _make_header_updater(i, "name"),
                            inputs=[name_tb, structured_rules_state, rules_version_state],
                            outputs=[structured_rules_state, rules_version_state],
                        )
                        enabled_cb.change(
                            _make_header_updater(i, "enabled"),
                            inputs=[enabled_cb, structured_rules_state, rules_version_state],
                            outputs=[structured_rules_state, rules_version_state],
                        )
                        remove_btn.click(
                            _make_remover(i),
                            inputs=[structured_rules_state, rules_version_state],
                            outputs=[structured_rules_state, rules_version_state],
                        )

            add_rule_btn.click(
                lambda rules, version: ((rules or []) + [new_rule()], version + 1),
                inputs=[structured_rules_state, rules_version_state],
                outputs=[structured_rules_state, rules_version_state],
            )

            def sync_rules_to_policy(rules, current_policy):
                """The rule list is the policy; render it for models that read text."""
                if rules_have_content(rules):
                    return render_rules_to_text(rules)
                return current_policy

            # Live limit warning + keep the effective policy text in sync.
            structured_rules_state.change(
                policy_limits_markdown,
                inputs=structured_rules_state,
                outputs=limits_display,
            )
            structured_rules_state.change(
                sync_rules_to_policy,
                inputs=[structured_rules_state, current_policy_state],
                outputs=current_policy_state,
            )
            # Driven off the rule list itself, so every route that changes the policy
            # (preset, upload, "Add as rules", a toggle, Clear) updates the same line.
            structured_rules_state.change(
                _policy_status,
                inputs=[structured_rules_state, policy_label_state],
                outputs=policy_status,
            )

        # Manual Edition Accordion
        with gr.Accordion("✏️ Manual Editing", open=False):
            manual_text = gr.Textbox(
                label="Policy Text",
                placeholder="Write or paste a policy, then add it as rules...",
                lines=20
            )
            add_as_rules_btn = gr.Button("➕ Add as rules", variant="primary")
            gr.Markdown(
                "Turns whatever is in the box into rules and appends them to the list "
                "above. A policy written as `## Name` with `Flag:` / `Don't flag:` / "
                "`Exception:` lines becomes one rule per section; anything else becomes "
                "a single rule you can then edit or delete."
            )
        
        policy_preview = gr.Markdown(value="*No policy loaded*")
        
        # Clear button
        clear_policy_btn = gr.Button("Clear Policy", variant="secondary")
        
        # Handlers
        def load_preset_handler(name, uploaded_policies, rules, version, label):
            """Append a preset's (or uploaded policy's) rules to the list.

            A preset used to replace the whole list and silently become the policy
            under test. Appending instead stacks it onto whatever is already there —
            each rule tagged with where it came from and switched on by default — so
            you can load several policies side by side and turn individual rules on or
            off with their checkbox instead of a load wiping your working set.
            """
            noop = (gr.update(), gr.update(), gr.update())
            if not name:
                return noop

            policy_text = ""
            if name in list_preset_policy_names():
                policy_text, _ = load_preset_policy(name, base_dir)
            elif name in (uploaded_policies or {}):
                policy_text = uploaded_policies[name]
            if not policy_text:
                return noop

            # Prefer per-section rules (they carry the definition context separately);
            # fall back to the catch-all split for a document with no "## " sections.
            parsed = parse_policy_markdown(policy_text)
            source = parsed or policy_text_to_categories(policy_text)
            added = [
                {"name": c.get("name", ""), "flag": c.get("violation_rule", ""),
                 "dont_flag": c.get("not_violation_rule", ""),
                 "exception": c.get("exception_override", ""),
                 "context": c.get("context", ""), "enabled": True, "source": name}
                for c in source
            ]
            if not added:
                return noop

            # Drop blank starter rows, keep anything with content, then append.
            kept = [r for r in (rules or []) if rule_has_content(r)]

            # Adding a preset that is already loaded is allowed — you may want a second
            # copy to edit against the first. But it is usually a double-click, and the
            # repeated rule names used to surface much later as a bare "duplicate
            # category title" from the PolicyLM endpoint, which named nothing the user
            # could act on. Say it here instead, where the cause is still on screen.
            present = {
                ((r.get("source") or "").strip(), (r.get("name") or "").strip().casefold())
                for r in kept
            }
            repeats = [
                a["name"] for a in added
                if (a["source"], (a["name"] or "").strip().casefold()) in present
            ]

            new_label = name if not kept else f"{(label or '').strip()} + {name}".lstrip(" +")
            if repeats:
                gr.Warning(
                    f"{name} is already in your policy — you now have two copies of "
                    f"{len(repeats)} rule{'' if len(repeats) == 1 else 's'} "
                    f"({', '.join(repeats[:3])}{'…' if len(repeats) > 3 else ''}). "
                    "Remove the duplicates with 🗑 unless you meant to compare two versions."
                )
            else:
                gr.Info(
                    f"Added {len(added)} rule{'' if len(added) == 1 else 's'} from {name} — "
                    f"your policy now has {len(kept) + len(added)}."
                )
            return kept + added, (version or 0) + 1, new_label

        load_preset_btn.click(
            load_preset_handler,
            inputs=[preset_dropdown, uploaded_policies_state,
                    structured_rules_state, rules_version_state, policy_label_state],
            outputs=[structured_rules_state, rules_version_state, policy_label_state],
        )

        def preview_preset(name, uploaded_policies):
            """Show which rules a policy contains before anyone adds it, so the button
            is concrete rather than a leap of faith."""
            if not name:
                return "*Select a policy above to preview its rules.*"
            if name in list_preset_policy_names():
                text, _ = load_preset_policy(name, base_dir)
            else:
                text = (uploaded_policies or {}).get(name, "")
            if not text:
                return "*Select a policy above to preview its rules.*"
            parsed = parse_policy_markdown(text) or policy_text_to_categories(text)
            names = [(r.get("name") or "Untitled").strip() for r in parsed]
            if not names:
                return "*This policy has no labelled sections; it will be added as a single rule.*"
            count = len(names)
            return f"Adds **{count} rule{'' if count == 1 else 's'}**: " + ", ".join(names)

        preset_dropdown.change(
            preview_preset,
            inputs=[preset_dropdown, uploaded_policies_state],
            outputs=preset_preview,
        )

        def load_upload_handler(f, uploaded_policies):
            """Make an uploaded policy available to add, without making it live.

            Uploading used to set the policy under test directly, so a file became the
            active policy the moment it landed. Now it only stores the file (so it
            appears in the "Select a policy to add" dropdown, preselected) and drops the
            text into the box for review — it goes live only when you add it as rules.
            """
            if not f:
                return gr.update(), "*No policy loaded*", gr.update(), {}

            filename = os.path.basename(f.name)
            upload_key = f"Uploaded - {filename}"
            policy_text, _ = load_policy_from_file(f.name)

            if uploaded_policies is None:
                uploaded_policies = {}
            is_duplicate = upload_key in uploaded_policies
            uploaded_policies[upload_key] = policy_text

            all_choices = list_preset_policy_names() + sorted(uploaded_policies.keys())
            if is_duplicate:
                gr.Warning(f"Policy '{filename}' already uploaded. Previous version overwritten.")

            return (
                policy_text,  # manual_text (for review; add it via the button)
                policy_text if policy_text else "*No policy loaded*",  # policy_preview
                gr.update(choices=all_choices, value=upload_key),  # preset_dropdown (preselected)
                uploaded_policies,  # uploaded_policies_state
            )

        upload_file.change(
            load_upload_handler,
            inputs=[upload_file, uploaded_policies_state],
            outputs=[manual_text, policy_preview, preset_dropdown, uploaded_policies_state],
        )
        
        def add_as_rules_handler(text, rules, version, label):
            """Append whatever is in the text box to the rule list.

            One button for both routes: a preset already arrives as rules, and this
            turns anything written or pasted into rules too, so the list is always the
            policy and nothing is silently overridden by something typed elsewhere.
            """
            if not (text or "").strip():
                return rules, version, label, gr.update()
            parsed = policy_text_to_categories(text)
            added = [
                {"name": c.get("name", ""), "flag": c.get("violation_rule", ""),
                 "dont_flag": c.get("not_violation_rule", ""),
                 "exception": c.get("exception_override", ""),
                 "context": c.get("context", ""), "enabled": True, "source": ""}
                for c in parsed
            ]
            kept = [r for r in (rules or []) if rule_has_content(r)]
            return kept + added, (version or 0) + 1, (label or "Custom (hand-written)"), gr.update(value="")

        def update_preview(text):
            return text if text else "*No policy loaded*"
        
        # Update state and preview when user leaves the textbox (blur event)
        add_as_rules_btn.click(
            add_as_rules_handler,
            inputs=[manual_text, structured_rules_state, rules_version_state, policy_label_state],
            outputs=[structured_rules_state, rules_version_state, policy_label_state, manual_text],
        )

        manual_text.blur(
            # Keeps the preview in step while writing. The text only becomes the policy
            # when "Add as rules" is pressed, so there is one route in and no silent
            # override of the rule list.
            lambda t: update_preview(t),
            inputs=manual_text,
            outputs=policy_preview,
        )
        
        clear_policy_btn.click(
            # Clear the rules too: clearing only the text left loaded rules live, so
            # the policy was not actually cleared.
            lambda v: ("", "", "*No policy loaded*", [new_rule()], (v or 0) + 1, ""),
            inputs=rules_version_state,
            outputs=[current_policy_state, manual_text, policy_preview,
                     structured_rules_state, rules_version_state, policy_label_state],
        )
        
        # Sync UI components when state changes externally (e.g., from dataset load)
        def sync_policy_ui(policy_text):
            preview_text = policy_text if policy_text else "*No policy loaded*"
            return policy_text, preview_text
        
        current_policy_state.change(
            sync_policy_ui,
            inputs=current_policy_state,
            outputs=[manual_text, policy_preview],
        )

    return {
        "current_policy_state": current_policy_state,
        "structured_rules_state": structured_rules_state,
        # Exposed so restoring a saved policy can force the rule editor to re-render.
        "rules_version_state": rules_version_state,
        "policy_label_state": policy_label_state,
    }