"""Tool for building and editing a literature-grounded report.

Intended to support the sqa task.

Uses the SQAResponse model to represent the report state.

Detects and removes orphaned citations when editing section text.

"""

from typing import Annotated, Any, Literal

from inspect_ai.tool import Tool, ToolDef, ToolError, ToolResult, tool
from inspect_ai.util import StoreModel, store_as
from pydantic import BaseModel, Field, RootModel
from shortuuid import uuid

from astabench.types.sqa import Citation, SQAResponse, SQASection

from .submission import get_submission_manager


class AddSectionsParams(BaseModel):
    action: Literal["add_sections"] = "add_sections"
    sections: list[SQASection] = Field(
        ...,
        description="List of new sections to append",
    )
    before_section: str | None = Field(
        None,
        description=(
            "Title of the section before which to insert new sections; "
            "if omitted, appends to end"
        ),
    )


class EditSectionParams(BaseModel):
    action: Literal["edit_section"] = "edit_section"
    target_section: str = Field(..., description="Title of the section to edit")
    new_title: str | None = Field(None, description="New title (optional)")
    text: str | None = Field(
        None,
        description="New full section text, or replacement string if 'old' is provided",
    )
    old: str | None = Field(
        None,
        description="Substring to replace within the section text; if provided, 'text' is treated as replacement",
    )
    citations: list[Citation] | None = Field(
        None,
        description="Optional list of Citation objects to replace this section's citations",
    )


class RemoveSectionsParams(BaseModel):
    action: Literal["remove_sections"] = "remove_sections"
    target_sections: list[str] = Field(
        ..., description="List of section titles to remove"
    )


class GetParams(BaseModel):
    action: Literal["get"] = "get"


class UndoParams(BaseModel):
    action: Literal["undo"] = "undo"


ReportEditorParams = RootModel[
    Annotated[
        AddSectionsParams
        | EditSectionParams
        | RemoveSectionsParams
        | GetParams
        | UndoParams,
        Field(discriminator="action"),
    ]
]


class ReportState(StoreModel):
    report: SQAResponse = Field(default_factory=lambda: SQAResponse())
    history: list[str] = Field(default_factory=list)


class ReportResponsePayload(BaseModel):
    report: SQAResponse
    orphaned_citations: list[Citation] = Field(default_factory=list)


DEFAULT_INSTANCE_ID = uuid()


@tool
def report_editor(
    instance: str | None = DEFAULT_INSTANCE_ID,
    history_limit: int = 3,
    enable_submission: bool = False,
) -> Tool:
    """Stateful tool to build and edit a SQAResponse report via a clear typed schema.

    Args:
        enable_submission: controls whether the "submit" action is available.
            Submissions are available via `astabench.tools.submission` with
            `has_submission` and `get_submission`.
    """

    async def execute(
        action: Literal[
            "add_sections",
            "edit_section",
            "remove_sections",
            "get",
            "undo",
            "submit",
        ],
        sections: list[SQASection] | None = None,
        before_section: str | None = None,
        remove_sections: list[str] | None = None,
        edit_params: EditSectionParams | None = None,
    ) -> ToolResult:
        """Execute report editing commands.

        Args:
          action: One of:
            - "add_sections": add new sections
            - "edit_section": update a section's title, text, or citations
            - "remove_sections": remove sections by title
            - "get": retrieve current report state
            - "undo": revert to previous state
            - "submit": submit the report as the final answer
          sections: For "add_sections", list of SQASection objects. Can be just section titles for outlining
          before_section: For "add_sections", title of the section to insert before (default to end)
          remove_sections: For "remove_sections", list of section titles to delete
          edit_params: Params for "edit_section"

        Returns:
          A JSON-serializable dict representing the updated report state and any orphaned citations.
        """
        state = store_as(ReportState, instance=instance)

        # prepare full rollback snapshot for mutating actions
        mutating = action not in ("undo", "get")
        if mutating:
            prior_history = state.history.copy()
            prior_report_json = state.report.model_dump_json()
            state.history.append(prior_report_json)
            # enforce maximum history length
            if len(state.history) > history_limit:
                state.history.pop(0)

        # Use a generic variable for params to satisfy mypy across cases
        p: Any

        # Track orphaned citations when section text is edited
        orphaned: list[Citation] = []

        # Validate params
        if action == "add_sections":
            if sections is None:
                raise ToolError("Missing sections for add_sections")
            p = AddSectionsParams(
                sections=sections,
                before_section=before_section,
            )
        elif action == "edit_section":
            if not edit_params:
                raise ToolError("Missing edit_params for edit_section")
            p = edit_params
        elif action == "remove_sections":
            if remove_sections is None:
                raise ToolError("Missing remove_sections for remove_sections")
            p = RemoveSectionsParams(target_sections=remove_sections)
        elif action == "get":
            p = GetParams()
        elif action == "undo":
            p = UndoParams()
        elif action == "submit":
            if not enable_submission:
                raise ValueError("`submit` action not enabled")
            if sections or before_section or remove_sections or edit_params:
                raise ToolError(
                    "Cannot submit report with pending changes. Please finalize edits and then use 'submit' with no other args."
                )
            if not state.report.sections:
                raise ToolError("Cannot submit empty report")
            submission = state.report.model_dump_json()
            get_submission_manager().write_submission(submission)
            return ReportResponsePayload(
                report=state.report,
                orphaned_citations=orphaned,
            ).model_dump_json()
        else:
            raise ToolError(f"Unknown action {action!r}")

        # Dispatch with automatic rollback on any exception
        try:
            match p:
                case AddSectionsParams():
                    new_secs = list(map(SQASection.model_validate, p.sections))
                    # ensure no duplicate titles in the bulk-add payload
                    titles = [sec.title for sec in new_secs]
                    if len(titles) != len(set(titles)):
                        raise ToolError("Duplicate titles in add_sections payload")
                    # prevent duplicate section titles
                    existing_titles = {s.title for s in state.report.sections}
                    for sec in new_secs:
                        if sec.title in existing_titles:
                            raise ToolError(
                                f"Duplicate section title: {sec.title}. If you wanted to edit this section, use the edit_section action instead."
                            )
                        existing_titles.add(sec.title)
                    if p.before_section is not None:
                        idx = next(
                            (
                                i
                                for i, s in enumerate(state.report.sections)
                                if s.title == p.before_section
                            ),
                            None,
                        )
                        if idx is None:
                            raise ToolError(
                                f"Section to insert before not found: {p.before_section}"
                            )
                        for offset, sec in enumerate(new_secs):
                            state.report.sections.insert(idx + offset, sec)
                    else:
                        state.report.sections.extend(new_secs)
                case EditSectionParams():
                    idx = next(
                        (
                            i
                            for i, s in enumerate(state.report.sections)
                            if s.title == p.target_section
                        ),
                        None,
                    )
                    if idx is None:
                        raise ToolError("Invalid section for edit")
                    sec = state.report.sections[idx]

                    # prevent renaming to an existing title
                    if p.new_title is not None:
                        other_titles = {
                            s.title
                            for s in state.report.sections
                            if s.title != p.target_section
                        }
                        if p.new_title in other_titles:
                            raise ToolError(f"Duplicate section title: {p.new_title}")
                        sec.title = p.new_title
                    if p.old is not None:
                        # perform substring replacement
                        sec_text: str = sec.text or ""
                        sec.text = sec_text.replace(p.old, p.text or "")
                    elif p.text is not None:
                        # replace entire section text
                        sec.text = p.text
                    if p.citations is not None:
                        # Merge: replace or append by id
                        merged: list[Citation] = []
                        existing_by_id = {c.id: c for c in sec.citations}
                        for new_cit in p.citations:
                            existing_by_id[new_cit.id] = new_cit
                        for c in sec.citations:
                            merged.append(existing_by_id.pop(c.id, c))
                        merged.extend(existing_by_id.values())
                        sec.citations = merged

                    # Detect orphaned citations and remove them
                    current_text = sec.text or ""
                    remaining: list[Citation] = []
                    for cit in sec.citations:
                        if cit.id in current_text:
                            remaining.append(cit)
                        else:
                            orphaned.append(cit)
                    sec.citations = remaining
                case RemoveSectionsParams():
                    # ensure all requested titles exist
                    existing = {s.title for s in state.report.sections}
                    for title in p.target_sections:
                        if title not in existing:
                            raise ToolError(f"Invalid section for removal: {title}")
                    # filter out all specified titles
                    state.report.sections = [
                        s
                        for s in state.report.sections
                        if s.title not in set(p.target_sections)
                    ]
                case GetParams():
                    pass
                case UndoParams():
                    if not state.history:
                        raise ToolError("No history to undo.")
                    last_state = state.history.pop()
                    state.report = SQAResponse.model_validate_json(last_state)
                case _:
                    raise ToolError(f"Unknown action {p.action!r}")
        except Exception:
            if mutating:
                state.history = prior_history
                state.report = SQAResponse.model_validate_json(prior_report_json)
            raise

        # Unified return payload
        payload = ReportResponsePayload(
            report=state.report,
            orphaned_citations=orphaned,
        )
        return payload.model_dump_json()

    tool_def = ToolDef(execute, name="report_editor")
    if not enable_submission:
        # Remove the 'submit' part of the description
        lines = tool_def.parameters.properties["action"].description.split("\n")
        tool_def.parameters.properties["action"].description = "\n".join(
            line for line in lines if not line.startswith('- "submit"')
        )
        tool_def.parameters.properties["action"].enum.remove("submit")

    return tool_def.as_tool()
