from __future__ import annotations

import json
import re
import secrets
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from app.settings import Settings


TASK_ID = re.compile(r"^cad_[a-z0-9]{12}$")
CONVERSATION_ID = re.compile(r"^conv_[a-z0-9]{12}$")
CANDIDATE_ID = re.compile(r"^candidate_[a-z0-9]{12,32}$")


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def new_id(prefix: str) -> str:
    return f"{prefix}_{secrets.token_hex(6)}"


def safe_task_id(task_id: str) -> str:
    value = str(task_id or "").strip()
    if not TASK_ID.fullmatch(value):
        raise ValueError("Invalid task id")
    return value


def safe_conversation_id(conversation_id: str) -> str:
    value = str(conversation_id or "").strip()
    if not CONVERSATION_ID.fullmatch(value):
        raise ValueError("Invalid conversation id")
    return value


def safe_relative_path(value: str) -> str:
    path = Path(str(value or ""))
    if not value or path.is_absolute() or ".." in path.parts:
        raise ValueError("Invalid artifact path")
    return path.as_posix()


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def read_json(path: Path, fallback: Any = None) -> Any:
    if not path.is_file():
        return fallback
    return json.loads(path.read_text(encoding="utf-8"))


class WorkspaceStore:
    def __init__(self, settings: Settings) -> None:
        self.settings = settings
        self.settings.task_root.mkdir(parents=True, exist_ok=True)
        self.settings.conversation_root.mkdir(parents=True, exist_ok=True)

    def task_dir(self, task_id: str) -> Path:
        return self.settings.task_root / safe_task_id(task_id)

    def task_path(self, task_id: str) -> Path:
        return self.task_dir(task_id) / "task.json"

    def conversation_dir(self, conversation_id: str) -> Path:
        return self.settings.conversation_root / safe_conversation_id(conversation_id)

    def conversation_path(self, conversation_id: str) -> Path:
        return self.conversation_dir(conversation_id) / "conversation.json"

    def write_tool_call_diagnostic(self, conversation_id: str, payload: dict[str, Any]) -> str:
        """Persist one failed model tool call without creating a CAD revision."""
        conversation = safe_conversation_id(conversation_id)
        relative = Path("diagnostics") / f"tool_call_{secrets.token_hex(8)}.json"
        path = self.conversation_dir(conversation) / relative
        write_json(path, payload)
        return (Path(conversation) / relative).as_posix()

    def write_cdsl_attempt(self, conversation_id: str, cdsl: Any, iteration: int) -> str:
        """Retain a parsed model candidate before validation or execution."""
        conversation = safe_conversation_id(conversation_id)
        relative = Path("diagnostics") / f"cdsl_attempt_{iteration:02d}_{secrets.token_hex(8)}.json"
        path = self.conversation_dir(conversation) / relative
        write_json(path, cdsl)
        return (Path(conversation) / relative).as_posix()

    def write_cdsl_validation_diagnostic(self, conversation_id: str, payload: dict[str, Any]) -> str:
        """Persist the reason a retained CDSL candidate was rejected."""
        conversation = safe_conversation_id(conversation_id)
        relative = Path("diagnostics") / f"cdsl_validation_{secrets.token_hex(8)}.json"
        path = self.conversation_dir(conversation) / relative
        write_json(path, payload)
        return (Path(conversation) / relative).as_posix()

    def ensure_conversation(
        self,
        conversation_id: str | None,
        current_task_id: str | None = None,
    ) -> dict[str, Any]:
        cid = safe_conversation_id(conversation_id) if conversation_id else new_id("conv")
        path = self.conversation_path(cid)
        current = read_json(path)
        if current:
            changed = False
            if current_task_id:
                current["current_task_id"] = safe_task_id(current_task_id)
                changed = True
            if changed:
                current["updated_at"] = now_iso()
                write_json(path, current)
            return current
        record = {
            "schema_version": "1.2",
            "conversation_id": cid,
            "created_at": now_iso(),
            "updated_at": now_iso(),
            "current_task_id": safe_task_id(current_task_id) if current_task_id else "",
            "messages": [],
            "attachments": [],
        }
        write_json(path, record)
        return record

    def read_conversation(self, conversation_id: str) -> dict[str, Any] | None:
        return read_json(self.conversation_path(conversation_id))

    def append_conversation_message(self, conversation_id: str, message: dict[str, Any], current_task_id: str | None = None) -> dict[str, Any]:
        record = self.ensure_conversation(conversation_id, current_task_id)
        known = {str(item.get("id")) for item in record["messages"]}
        if str(message.get("id")) not in known:
            record["messages"].append(message)
        if current_task_id:
            record["current_task_id"] = safe_task_id(current_task_id)
        record["updated_at"] = now_iso()
        write_json(self.conversation_path(record["conversation_id"]), record)
        return record

    def write_conversation_upload(self, conversation_id: str, filename: str, data: bytes) -> tuple[str, Path]:
        conversation = safe_conversation_id(conversation_id)
        safe_name = re.sub(r"[^a-zA-Z0-9._-]+", "_", Path(filename).name).strip("._") or "attachment"
        relative = Path("uploads") / f"upload_{secrets.token_hex(6)}_{safe_name}"
        target = self.conversation_attachment_path(conversation, relative.as_posix())
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_bytes(data)
        return relative.as_posix(), target

    def add_conversation_attachment(self, conversation_id: str, attachment: dict[str, Any]) -> dict[str, Any]:
        conversation = safe_conversation_id(conversation_id)
        record = self.read_conversation(conversation)
        if record is None:
            raise ValueError("Conversation not found")
        if str(attachment.get("conversation_id") or "") != conversation:
            raise ValueError("Attachment does not belong to this conversation")
        attachment_id = str(attachment.get("id") or "")
        if not attachment_id:
            raise ValueError("Attachment id is required")
        self.conversation_attachment_path(conversation, str(attachment.get("path") or ""))
        attachments = record.setdefault("attachments", [])
        if any(str(item.get("id") or "") == attachment_id for item in attachments if isinstance(item, dict)):
            raise ValueError("Attachment already exists")
        attachments.append(attachment)
        record["updated_at"] = now_iso()
        write_json(self.conversation_path(conversation), record)
        return record

    def ensure_task(self, task_id: str | None, request: str) -> dict[str, Any]:
        tid = safe_task_id(task_id) if task_id else new_id("cad")
        path = self.task_path(tid)
        current = read_json(path)
        if current:
            return self._migrate_task(current, path)
        task_dir = self.task_dir(tid)
        (task_dir / "revisions").mkdir(parents=True, exist_ok=True)
        source_requirements = str(request or "").strip()
        source_path = task_dir / "source-requirements.md"
        source_path.write_text(source_requirements + "\n", encoding="utf-8")
        record = {
            "schema_version": "2.0",
            "task_id": tid,
            "request": request,
            "source_requirements_path": "source-requirements.md",
            "created_at": now_iso(),
            "updated_at": now_iso(),
            "current_revision": "",
            "active_revision": "",
            "published_revision": "",
            "lifecycle": "completed",
            "run_id": "",
            "requirements_path": "",
            "agent_state_path": "",
            "active_candidate_id": "",
            "active_branch_id": "main",
            "run_failure_path": "",
            "revisions": [],
        }
        write_json(path, record)
        return record

    def _migrate_task(self, task: dict[str, Any], path: Path) -> dict[str, Any]:
        """Fill in omitted fields for native autonomous tasks only.

        A task from a retired protocol must never be made to look resumable by
        rewriting its schema version.  In particular, doing so would make an
        old revision and its obsolete planning state appear to be a valid
        autonomous work head.  Historical tasks remain read-only records; the
        service startup path marks only *running* legacy tasks as obsolete.
        """
        if str(task.get("schema_version") or "") != "2.0":
            return task
        changed = False
        current = str(task.get("current_revision") or "")
        defaults = {
            "active_revision": current,
            "published_revision": current,
            "lifecycle": "completed",
            "run_id": "",
            "requirements_path": "",
            "source_requirements_path": "",
            "completion_checklist_path": "",
            "agent_state_path": "",
            "active_candidate_id": "",
            "active_branch_id": "main",
            "run_failure_path": "",
        }
        for key, value in defaults.items():
            if key not in task:
                task[key] = value
                changed = True
        if not str(task.get("source_requirements_path") or ""):
            relative = Path("source-requirements.md")
            source_path = path.parent / relative
            if not source_path.exists():
                source_path.write_text(str(task.get("request") or "").strip() + "\n", encoding="utf-8")
            task["source_requirements_path"] = relative.as_posix()
            changed = True
        for revision in task.get("revisions") or ():
            if not isinstance(revision, dict):
                continue
            if "visibility" not in revision:
                revision["visibility"] = "final" if str(revision.get("revision_id") or "") == str(task["published_revision"] or "") else "checkpoint"
                changed = True
            if "branch_id" not in revision:
                revision["branch_id"] = "main"
                changed = True
        if changed:
            task["updated_at"] = now_iso()
            write_json(path, task)
        return task

    def next_revision(self, task_id: str) -> tuple[str, Path]:
        task = self.ensure_task(task_id, "")
        revision_id = f"rev_{len(task['revisions']) + 1:03d}"
        revision_dir = self.task_dir(task_id) / "revisions" / revision_id
        revision_dir.mkdir(parents=True, exist_ok=False)
        return revision_id, revision_dir

    def update_task(self, task_id: str, revision: dict[str, Any]) -> dict[str, Any]:
        task = self.ensure_task(task_id, "")
        task["revisions"].append(revision)
        if revision.get("status") == "success":
            task["current_revision"] = revision["revision_id"]
            task["active_revision"] = revision["revision_id"]
            if revision.get("visibility") == "final":
                task["published_revision"] = revision["revision_id"]
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return task

    def read_task(self, task_id: str) -> dict[str, Any] | None:
        task = read_json(self.task_path(task_id))
        return self._migrate_task(task, self.task_path(task_id)) if isinstance(task, dict) else None

    def start_generation(self, task_id: str, *, request: str, run_id: str | None = None) -> dict[str, Any]:
        task = self.ensure_task(task_id, request)
        if str(task.get("lifecycle") or "") == "running":
            raise ValueError("CAD task is already running")
        task.update({
            "lifecycle": "running",
            "run_id": run_id or new_id("run"),
            "active_candidate_id": "",
            "run_failure_path": "",
            "request": request or task.get("request") or "",
            "active_revision": str(task.get("current_revision") or ""),
            "active_branch_id": str(task.get("active_branch_id") or "main"),
            "updated_at": now_iso(),
        })
        write_json(self.task_path(task_id), task)
        return task

    def finish_generation(self, task_id: str, *, lifecycle: str, failure: dict[str, Any] | None = None) -> dict[str, Any]:
        if lifecycle not in {"completed", "failed"}:
            raise ValueError("Generation lifecycle must be completed or failed")
        task = self.ensure_task(task_id, "")
        failure_path = ""
        if failure:
            failure_path = "run-failures/" + f"failure_{secrets.token_hex(8)}.json"
            write_json(self.task_dir(task_id) / failure_path, failure)
        if lifecycle == "completed":
            task["published_revision"] = str(task.get("active_revision") or task.get("current_revision") or "")
            for revision in task.get("revisions") or ():
                if isinstance(revision, dict) and revision.get("revision_id") == task["published_revision"]:
                    revision["visibility"] = "final"
        task.update({
            "lifecycle": lifecycle,
            "active_candidate_id": "",
            "run_failure_path": failure_path,
            "updated_at": now_iso(),
        })
        write_json(self.task_path(task_id), task)
        return task

    def set_active_revision(self, task_id: str, revision_id: str, *, branch_id: str | None = None) -> dict[str, Any]:
        task = self.ensure_task(task_id, "")
        if not revision_id:
            task["active_revision"] = ""
            task["current_revision"] = ""
            if branch_id:
                task["active_branch_id"] = branch_id
            task["updated_at"] = now_iso()
            write_json(self.task_path(task_id), task)
            return task
        revision = next((item for item in task.get("revisions") or () if isinstance(item, dict) and item.get("revision_id") == revision_id), None)
        if not isinstance(revision, dict) or revision.get("status") != "success":
            raise ValueError("Active revision must be a successful revision")
        task["active_revision"] = revision_id
        task["current_revision"] = revision_id
        if branch_id:
            task["active_branch_id"] = branch_id
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return task

    def set_agent_state(self, task_id: str, values: dict[str, Any]) -> dict[str, Any]:
        """Persist enough autonomous-agent state to resume after a restart."""
        task = self.ensure_task(task_id, "")
        relative = Path("agent-state.json")
        write_json(self.task_dir(task_id) / relative, values)
        task["agent_state_path"] = relative.as_posix()
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return task

    def read_agent_state(self, task_id: str) -> dict[str, Any] | None:
        task = self.read_task(task_id) or {}
        relative = str(task.get("agent_state_path") or "")
        value = read_json(self.artifact_path(task_id, relative)) if relative else None
        return value if isinstance(value, dict) else None

    def requirements_document_path(self, task_id: str) -> Path:
        return self.task_dir(task_id) / "requirements.md"

    def read_source_requirements(self, task_id: str) -> str:
        """Return the server-owned, immutable user request for this task.

        Older tasks predate the artifact. Their persisted request is a
        read-only compatibility fallback; new tasks always have the artifact.
        """
        task = self.read_task(task_id) or {}
        relative = str(task.get("source_requirements_path") or "")
        if relative:
            path = self.artifact_path(task_id, relative)
            if path.is_file():
                return path.read_text(encoding="utf-8")
        return str(task.get("request") or "")

    def write_requirements_document(self, task_id: str, markdown: str) -> Path:
        """Write the one immutable requirements document for an agent run."""
        text = str(markdown or "").strip()
        if not text:
            raise ValueError("requirements.md must not be empty")
        task = self.ensure_task(task_id, "")
        relative = Path("requirements.md")
        path = self.task_dir(task_id) / relative
        if path.exists() or str(task.get("requirements_path") or ""):
            raise ValueError("requirements.md is frozen and cannot be rewritten")
        path.write_text(text + "\n", encoding="utf-8")
        task["requirements_path"] = relative.as_posix()
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return path

    def read_requirements_document(self, task_id: str) -> str:
        task = self.read_task(task_id) or {}
        relative = str(task.get("requirements_path") or "")
        if not relative:
            return ""
        path = self.artifact_path(task_id, relative)
        return path.read_text(encoding="utf-8") if path.is_file() else ""

    def completion_checklist_path(self, task_id: str) -> Path:
        return self.task_dir(task_id) / "completion.md"

    def write_completion_checklist(self, task_id: str, markdown: str) -> Path:
        """Persist the one immutable, author-owned completion checklist."""
        text = str(markdown or "").strip()
        if not text:
            raise ValueError("completion.md must not be empty")
        task = self.ensure_task(task_id, "")
        if not self.read_requirements_document(task_id):
            raise ValueError("requirements.md must be written before completion.md")
        relative = Path("completion.md")
        path = self.task_dir(task_id) / relative
        if path.exists() or str(task.get("completion_checklist_path") or ""):
            raise ValueError("completion.md is frozen and cannot be rewritten")
        path.write_text(text + "\n", encoding="utf-8")
        task["completion_checklist_path"] = relative.as_posix()
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return path

    def read_completion_checklist(self, task_id: str) -> str:
        task = self.read_task(task_id) or {}
        relative = str(task.get("completion_checklist_path") or "")
        if not relative:
            return ""
        path = self.artifact_path(task_id, relative)
        return path.read_text(encoding="utf-8") if path.is_file() else ""

    def new_candidate(self, task_id: str) -> tuple[str, Path]:
        task = self.ensure_task(task_id, "")
        candidate_id = f"candidate_{secrets.token_hex(8)}"
        path = self.task_dir(task_id) / "candidates" / candidate_id
        path.mkdir(parents=True, exist_ok=False)
        task["active_candidate_id"] = candidate_id
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return candidate_id, path

    def candidate_dir(self, task_id: str, candidate_id: str) -> Path:
        if not CANDIDATE_ID.fullmatch(str(candidate_id or "")):
            raise ValueError("Invalid candidate id")
        return self.task_dir(task_id) / "candidates" / candidate_id

    def clear_active_candidate(self, task_id: str, candidate_id: str | None = None) -> dict[str, Any]:
        task = self.ensure_task(task_id, "")
        if candidate_id and str(task.get("active_candidate_id") or "") not in {"", candidate_id}:
            raise ValueError("Candidate is not active")
        task["active_candidate_id"] = ""
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return task

    def append_agent_audit(self, task_id: str, kind: str, payload: dict[str, Any]) -> str:
        safe_kind = re.sub(r"[^a-z0-9_-]+", "-", str(kind).lower()).strip("-") or "event"
        relative = Path("agent-audit") / f"{now_iso().replace(':', '-').replace('+', '_')}_{safe_kind}_{secrets.token_hex(4)}.json"
        write_json(self.task_dir(task_id) / relative, {"recorded_at": now_iso(), "kind": safe_kind, **payload})
        return relative.as_posix()

    def update_revision_metadata(self, task_id: str, revision_id: str, values: dict[str, Any]) -> dict[str, Any]:
        task = self.ensure_task(task_id, "")
        revision = next((item for item in task.get("revisions") or () if isinstance(item, dict) and item.get("revision_id") == revision_id), None)
        if not isinstance(revision, dict):
            raise ValueError("Revision does not exist")
        revision.update(values)
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return task

    def rollback_to_revision(self, task_id: str, revision_id: str, *, branch_id: str) -> dict[str, Any]:
        """Move the generation head without deleting immutable checkpoint artifacts."""
        task = self.ensure_task(task_id, "")
        source_branch_id = str(task.get("active_branch_id") or "main")
        if revision_id:
            target = next(
                (item for item in task.get("revisions") or () if isinstance(item, dict) and item.get("revision_id") == revision_id),
                None,
            )
            if not isinstance(target, dict) or target.get("status") != "success":
                raise ValueError("Active revision must be a successful revision")
        children: dict[str, set[str]] = {}
        for revision in task.get("revisions") or ():
            if not isinstance(revision, dict):
                continue
            parent = str(revision.get("parent_revision_id") or "")
            child = str(revision.get("revision_id") or "")
            if child and str(revision.get("branch_id") or "main") == source_branch_id:
                children.setdefault(parent, set()).add(child)
        superseded: set[str] = set()
        pending = list(children.get(revision_id, set()))
        while pending:
            child = pending.pop()
            if not child or child in superseded:
                continue
            superseded.add(child)
            pending.extend(children.get(child, set()))
        for revision in task.get("revisions") or ():
            if isinstance(revision, dict) and str(revision.get("revision_id") or "") in superseded and revision.get("visibility") == "checkpoint":
                revision["visibility"] = "superseded"
        task["active_revision"] = revision_id
        task["current_revision"] = revision_id
        task["active_branch_id"] = branch_id
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return task

    def rollback_anchor_for_nodes(
        self,
        task_id: str,
        node_ids: list[str],
        *,
        fallback_revision_id: str = "",
    ) -> str:
        """Return the revision before every affected node's latest checkpoint.

        Returning each affected revision's parent (rather than the revision
        itself) ensures the faulty node is regenerated.  A common ancestor
        keeps unrelated upstream work intact while permitting a single rollback
        over any number of affected nodes.
        """
        task = self.read_task(task_id) or {}
        revisions = [item for item in task.get("revisions") or () if isinstance(item, dict)]
        by_id = {str(item.get("revision_id") or ""): item for item in revisions}
        parents: list[str] = []
        for node_id in dict.fromkeys(str(item) for item in node_ids if str(item)):
            matching = [item for item in revisions if item.get("status") == "success" and str(item.get("node_id") or "") == node_id]
            if matching:
                parents.append(str(matching[-1].get("parent_revision_id") or ""))
        if not parents:
            return fallback_revision_id

        def lineage(revision_id: str) -> list[str]:
            chain = [revision_id]
            seen = {revision_id}
            current = revision_id
            while current:
                parent = str((by_id.get(current) or {}).get("parent_revision_id") or "")
                if parent in seen:
                    break
                chain.append(parent)
                seen.add(parent)
                current = parent
            return chain

        common = set(lineage(parents[0]))
        for parent in parents[1:]:
            common.intersection_update(lineage(parent))
        if not common:
            return fallback_revision_id
        return next((revision for revision in lineage(parents[0]) if revision in common), fallback_revision_id)

    def write_generation_failure(self, task_id: str, payload: dict[str, Any]) -> str:
        """Persist an attempt-level diagnostic without changing lifecycle."""
        relative = Path("generation-failures") / f"failure_{secrets.token_hex(8)}.json"
        write_json(self.task_dir(task_id) / relative, payload)
        return relative.as_posix()

    def running_tasks(self) -> list[dict[str, Any]]:
        """Enumerate durable tasks that need a process-local worker."""
        tasks: list[dict[str, Any]] = []
        for candidate in self.settings.task_root.glob("cad_*"):
            if not candidate.is_dir() or not TASK_ID.fullmatch(candidate.name):
                continue
            task = self.read_task(candidate.name)
            if isinstance(task, dict) and task.get("lifecycle") == "running":
                tasks.append(task)
        return tasks

    def revision_dir(self, task_id: str, revision_id: str) -> Path:
        return self.task_dir(task_id) / "revisions" / revision_id

    def current_cdsl_path(self, task_id: str) -> Path | None:
        task = self.read_task(task_id)
        revision_id = str((task or {}).get("active_revision") or (task or {}).get("current_revision") or "")
        if not revision_id:
            return None
        candidate = self.task_dir(task_id) / "revisions" / revision_id / "model.cdsl.json"
        return candidate if candidate.is_file() else None

    def revision_topology_path(self, task_id: str, revision_id: str) -> Path | None:
        task = self.read_task(task_id)
        revision = next(
            (item for item in (task or {}).get("revisions") or [] if str(item.get("revision_id") or "") == revision_id),
            None,
        )
        if not isinstance(revision, dict):
            return None
        relative = str(revision.get("topology_path") or "")
        if not relative:
            return None
        candidate = self.artifact_path(task_id, relative)
        return candidate if candidate.is_file() else None

    def current_topology_path(self, task_id: str) -> Path | None:
        task = self.read_task(task_id)
        revision_id = str((task or {}).get("active_revision") or (task or {}).get("current_revision") or "")
        if not revision_id:
            return None
        return self.revision_topology_path(task_id, revision_id)

    def revision_cdsl_path(self, task_id: str, revision_id: str) -> Path | None:
        task = self.read_task(task_id)
        revision = next(
            (item for item in (task or {}).get("revisions") or [] if str(item.get("revision_id") or "") == revision_id),
            None,
        )
        if not isinstance(revision, dict):
            return None
        relative = str(revision.get("cdsl_path") or "")
        if not relative:
            return None
        candidate = self.artifact_path(task_id, relative)
        return candidate if candidate.is_file() else None

    def artifact_path(self, task_id: str, relative_path: str) -> Path:
        safe = safe_relative_path(relative_path)
        root = self.task_dir(task_id).resolve()
        target = (root / safe).resolve()
        if root != target and root not in target.parents:
            raise ValueError("Artifact path escapes task directory")
        return target

    def conversation_attachment_path(self, conversation_id: str, relative_path: str) -> Path:
        safe = safe_relative_path(relative_path)
        root = self.conversation_dir(conversation_id).resolve()
        target = (root / safe).resolve()
        if root != target and root not in target.parents:
            raise ValueError("Attachment path escapes conversation directory")
        return target