"""The run-scoped work branch: durable, resumable codegen workspace.

Three branch roles (see docs/CODEGEN_WORK_BRANCH.md):

    main                     production / source truth
    preview                  the currently previewable, deployable project state
    p2hs/codegen/<run_id>    one isolated, resumable workspace per run

The runner works on the run branch and commits/pushes checkpoints to it; ``preview`` is only
advanced (by fast-forward or a clean merge) after Core-owned validation passes, so a failed
run can never destroy the last working preview and a retry resumes from the branch instead of
regenerating everything. Rules this module encodes:

- created ONCE per run (from the preview head at that moment), never once per retry;
- never force-pushed: commits are the checkpoints, the database only remembers where they are;
- retention: failed/blocked branches are kept for investigation and retry, then cleaned up.

State lives in ``plan.context.work_branch`` (not ``checkpoint_payload``: retry clears that),
which is also what dispatch spreads into the runner payload.
"""

from __future__ import annotations

import re
from datetime import UTC, datetime, timedelta
from typing import Any
from uuid import UUID

import structlog

logger = structlog.get_logger()

WORK_BRANCH_PREFIX = "p2hs/codegen/"
_BRANCH_RE = re.compile(r"^p2hs/codegen/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$")
_SHA_RE = re.compile(r"^[0-9a-f]{40}$")

# Phases the runner can prove it reached, in order.
# PARTIAL_WORK is work saved by a run that failed before the agent finished: a retry must run the
# agent again on top of it, never skip the agent.
PARTIAL_WORK = "partial_work"
AGENT_OUTPUT = "agent_output"
CAPABILITIES_PROVISIONED = "capabilities_provisioned"
PROMOTED = "promoted"
CHECKPOINTS = (PARTIAL_WORK, AGENT_OUTPUT, CAPABILITIES_PROVISIONED, PROMOTED)

# The failure class the coding agent itself can fix (a malformed capability declaration).
CORRECTABLE_CLASS = "invalid_requirement"
_MAX_MESSAGE = 500
_MAX_ERRORS = 20


def work_branch_name(run_id: UUID | str) -> str:
    return f"{WORK_BRANCH_PREFIX}{run_id}"


def is_work_branch(name: str) -> bool:
    return bool(_BRANCH_RE.fullmatch(name or ""))


def is_sha(value: str | None) -> bool:
    return bool(_SHA_RE.fullmatch(value or ""))


def _now() -> str:
    return datetime.now(UTC).isoformat()


def new_state(run_id: UUID | str, base_branch: str, base_sha: str) -> dict[str, Any]:
    """State for a branch that was just created from ``base_branch`` at ``base_sha``."""
    return {
        "branch": work_branch_name(run_id),
        "base_branch": base_branch,
        "base_sha": base_sha,
        "last_checkpoint": None,
        "last_checkpoint_sha": None,
        "status": "active",
        "retryable": False,
        "failure_class": None,
        "failure_message": None,
        "correction": None,
        "correction_attempts": 0,
        "updated_at": _now(),
    }


def valid_state(state: Any, run_id: UUID | str | None = None) -> bool:
    """A stored state is only trusted if it names this run's branch and a real base commit."""
    if not isinstance(state, dict) or not is_work_branch(str(state.get("branch") or "")):
        return False
    if run_id is not None and state["branch"] != work_branch_name(run_id):
        return False
    return is_sha(state.get("base_sha"))


def apply_report(
    state: dict[str, Any],
    *,
    status: str,
    checkpoint: str | None = None,
    checkpoint_sha: str | None = None,
    failure_class: str | None = None,
    retryable: bool | None = None,
    message: str | None = None,
    correction: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Fold one runner/Core report into the state and return the new state.

    ``checkpoint`` only ever moves forward (a stale or repeated callback cannot rewind it)."""
    new = dict(state)
    if checkpoint in CHECKPOINTS and is_sha(checkpoint_sha):
        current = state.get("last_checkpoint")
        if current not in CHECKPOINTS or CHECKPOINTS.index(checkpoint) >= CHECKPOINTS.index(current):
            new["last_checkpoint"] = checkpoint
            new["last_checkpoint_sha"] = checkpoint_sha
    new["status"] = status
    if status in {"failed", "blocked"}:
        new["failure_class"] = failure_class
        new["retryable"] = bool(retryable)
        new["failure_message"] = (message or "")[:_MAX_MESSAGE] or None
        new["correction"] = bounded_correction(correction)
    else:
        new["failure_class"] = None
        new["retryable"] = False
        new["failure_message"] = None
        new["correction"] = None
    new["updated_at"] = _now()
    return new


def bounded_correction(correction: Any) -> dict[str, Any] | None:
    """Keep only the structured, bounded part of a correction request (paths, codes and
    messages of our own validation; renames). Anything else from a runner is dropped."""
    if not isinstance(correction, dict):
        return None
    entries = []
    for entry in (correction.get("entries") or [])[:_MAX_ERRORS]:
        if not isinstance(entry, dict):
            continue
        entries.append(
            {
                "name": str(entry.get("name") or "")[:100],
                "errors": [
                    {
                        "path": str(e.get("path") or "")[:200],
                        "code": str(e.get("code") or "")[:60],
                        "message": str(e.get("message") or "")[:300],
                    }
                    for e in (entry.get("errors") or [])[:_MAX_ERRORS]
                    if isinstance(e, dict)
                ],
                "renamed": [
                    {
                        "path": str(r.get("path") or "")[:200],
                        "from": str(r.get("from") or "")[:200],
                        "to": str(r.get("to") or "")[:200],
                    }
                    for r in (entry.get("renamed") or [])[:_MAX_ERRORS]
                    if isinstance(r, dict)
                ],
            }
        )
    return {"entries": entries} if entries else None


def resume_plan(state: Any, *, max_corrections: int) -> dict[str, Any] | None:
    """What a retry of this run should do, or None to just run normally on the branch.

    - A correctable failure (the agent's own manifest was invalid) re-runs the agent on top of
      its existing work with the exact errors, up to ``max_corrections`` times.
    - Any other failure at or after the agent finished (platform bug, Cloud outage, a capability
      the owner has since configured, a diverged preview) only needs the gate and promotion
      redone: the agent is skipped and the frontend is not regenerated.
    """
    if not valid_state(state) or state.get("last_checkpoint") not in CHECKPOINTS:
        return None
    if state.get("last_checkpoint") == PARTIAL_WORK:
        return None  # the agent never finished: run it again, on top of the saved work
    if state.get("status") not in {"failed", "blocked"}:
        return None
    if state.get("failure_class") == CORRECTABLE_CLASS:
        correction = state.get("correction")
        if correction and int(state.get("correction_attempts") or 0) < max_corrections:
            return {"skip_agent": False, "correction": correction,
                    "last_checkpoint": state["last_checkpoint"]}
        return None
    return {"skip_agent": True, "last_checkpoint": state["last_checkpoint"]}


def begin_retry(state: dict[str, Any], resume: dict[str, Any] | None) -> tuple[dict[str, Any], dict[str, Any]]:
    """State and plan-context additions for a retry that follows ``resume_plan``.

    The branch, its base and its checkpoint are kept (the branch is never recreated); the failure is
    cleared and the run is active again. A correction attempt is counted the moment it starts.
    Returns (new_state, context_updates) where context_updates holds ``resume``/``correction``."""
    new = dict(state)
    new["status"] = "active"
    new["failure_class"] = None
    new["retryable"] = False
    new["failure_message"] = None
    new["correction"] = None
    new["updated_at"] = _now()
    updates: dict[str, Any] = {}
    if resume:
        updates["resume"] = {
            "skip_agent": bool(resume.get("skip_agent")),
            "last_checkpoint": resume.get("last_checkpoint"),
        }
        if resume.get("correction"):
            new["correction_attempts"] = int(state.get("correction_attempts") or 0) + 1
            updates["correction"] = {**resume["correction"], "attempt": new["correction_attempts"]}
    return new, updates


def build_correction_appendix(correction: dict[str, Any], attempt: int) -> str:
    """Prompt text for a correction attempt: the agent's own work is already on the branch."""
    is_auth_conformance = any(
        str(entry.get("name") or "") == "generated-auth-conformance"
        for entry in correction.get("entries") or []
        if isinstance(entry, dict)
    )
    if is_auth_conformance:
        lines = [
            f"## Security correction needed (attempt {attempt})",
            "The platform's generated-source auth conformance gate found direct auth transport calls.",
            "Fix the authored UI/components in this existing repository; do not change the platform",
            "helper, locked files, manifests, or unrelated product behavior.",
            "Replace every direct `/api/auth/*`, `/v1/projects/*/auth/*`, or `auth/me` request with",
            "the exact canonical exports from the installed helper: `me`, `signIn`, `signUp`,",
            "`signOut`, `updateProfile`, or `requestMagicLink` as appropriate. Read the helper and",
            "the local P2HS recipe first. Do not copy its implementation into generated source.",
            "After editing, search authored source again and commit the correction.",
        ]
        for entry in correction.get("entries") or []:
            if not isinstance(entry, dict):
                continue
            for error in entry.get("errors") or []:
                lines.append(f"- `{error.get('path') or '(generated source)'}`: {error.get('message') or 'direct auth transport'} ({error.get('code') or 'AUTH-GEN-007'})")
        return "\n".join(lines)
    lines = [
        f"## Correction needed (attempt {attempt})",
        "Your work from the previous attempt is already in this repository. The platform "
        "validated the capability declaration in `p2hs/feature-packs.json` and it could not be "
        "accepted. Fix ONLY the problems below in that file. Do not regenerate or restructure "
        "the rest of the app. The platform will not guess anything security- or "
        "data-relevant for you: an unrecognized key, action, actor, workflow transition or "
        "field type must be corrected by you, never worked around.",
    ]
    renamed = False
    for entry in correction.get("entries") or []:
        lines.append(f"\n### `{entry.get('name') or 'capability'}`")
        for error in entry.get("errors") or []:
            lines.append(f"- `{error['path'] or '(entry)'}`: {error['message']} ({error['code']})")
        for rename in entry.get("renamed") or []:
            renamed = True
            lines.append(
                f"- The platform normalized `{rename['path']}` from `{rename['from']}` to "
                f"`{rename['to']}`; use the canonical name."
            )
    if renamed:
        lines.append(
            "\nWherever the app code refers to a renamed entity, field or resource (forms, API "
            "calls, records), change it to the canonical name."
        )
    return "\n".join(lines)


async def persist_state(db_factory, run_id: UUID, state: dict[str, Any]) -> bool:
    """Store the state on the run's plan (read-modify-write). Best-effort: returns False and
    logs instead of raising, because losing the pointer must not fail a run (the branch and
    its commits are the source of truth)."""
    from src.domain.codegen.repository import CodegenRunRepository

    try:
        async with db_factory() as db:
            run = await CodegenRunRepository(db).get_by_id(run_id)
            if not run:
                return False
            plan = dict(run.plan or {})
            context = dict(plan.get("context") or {})
            context["work_branch"] = state
            plan["context"] = context
            run.plan = plan
            await db.commit()
        return True
    except Exception as exc:
        logger.warning(
            "work_branch_state_persist_failed", run_id=str(run_id), error_type=type(exc).__name__
        )
        return False


def is_expired(state: Any, now: datetime, retention_days: int) -> bool:
    """True when a branch is old enough to delete: terminal state kept for ``retention_days``
    after its last update (promoted branches are safe to drop much sooner, see the caller)."""
    if not isinstance(state, dict) or state.get("status") == "active":
        return False
    try:
        updated = datetime.fromisoformat(str(state.get("updated_at")))
    except (TypeError, ValueError):
        return False
    if updated.tzinfo is None:
        updated = updated.replace(tzinfo=UTC)
    return now - updated >= timedelta(days=retention_days)