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}$")


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 write_conversation_planning(self, conversation_id: str, prefix: str, payload: dict[str, Any]) -> str:
        """Persist structured intake/planning evidence before a task exists."""
        conversation = safe_conversation_id(conversation_id)
        safe_prefix = re.sub(r"[^a-zA-Z0-9_-]+", "-", prefix).strip("-") or "planning"
        relative = Path("planning") / f"{safe_prefix}-{secrets.token_hex(6)}.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)
        record = {
            "schema_version": "1.3",
            "task_id": tid,
            "request": request,
            "created_at": now_iso(),
            "updated_at": now_iso(),
            "current_revision": "",
            "active_revision": "",
            "published_revision": "",
            "lifecycle": "completed",
            "run_id": "",
            "generation_spec_path": "",
            "run_context_path": "",
            "active_node_id": "",
            "run_failure_path": "",
            "revisions": [],
        }
        write_json(path, record)
        return record

    def _migrate_task(self, task: dict[str, Any], path: Path) -> dict[str, Any]:
        """Add run-state fields lazily without rewriting successful history."""
        changed = False
        current = str(task.get("current_revision") or "")
        defaults = {
            "schema_version": "1.3",
            "active_revision": current,
            "published_revision": current,
            "lifecycle": "completed",
            "run_id": "",
            "generation_spec_path": "",
            "run_context_path": "",
            "active_node_id": "",
            "run_failure_path": "",
        }
        for key, value in defaults.items():
            if key not in task:
                task[key] = value
                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_node_id": "",
            "run_failure_path": "",
            "request": request or task.get("request") or "",
            "active_revision": str(task.get("current_revision") or ""),
            "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_node_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"] = ""
            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_active_node(self, task_id: str, node_id: str) -> dict[str, Any]:
        task = self.ensure_task(task_id, "")
        task["active_node_id"] = node_id
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return task

    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.set_active_revision(task_id, revision_id, branch_id=branch_id)
        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 parent and child:
                children.setdefault(parent, set()).add(child)
        superseded: set[str] = set()
        pending = list(children.get(revision_id, set())) if revision_id else [
            str(item.get("revision_id") or "")
            for item in task.get("revisions") or ()
            if isinstance(item, dict) and not str(item.get("parent_revision_id") or "")
        ]
        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["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 generation_spec_path(self, task_id: str) -> Path:
        return self.task_dir(task_id) / "generation-spec.json"

    def write_generation_spec(self, task_id: str, spec: dict[str, Any]) -> Path:
        task = self.ensure_task(task_id, "")
        path = self.generation_spec_path(task_id)
        write_json(path, spec)
        task["generation_spec_path"] = path.relative_to(self.task_dir(task_id)).as_posix()
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return path

    def generation_run_context_path(self, task_id: str) -> Path:
        return self.task_dir(task_id) / "generation-run-context.json"

    def write_generation_run_context(self, task_id: str, context: dict[str, Any]) -> Path:
        """Persist the frozen authoring inputs needed to resume after restart."""
        task = self.ensure_task(task_id, "")
        path = self.generation_run_context_path(task_id)
        write_json(path, context)
        task["run_context_path"] = path.relative_to(self.task_dir(task_id)).as_posix()
        task["updated_at"] = now_iso()
        write_json(self.task_path(task_id), task)
        return path

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

    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 read_generation_spec(self, task_id: str) -> dict[str, Any] | None:
        task = self.read_task(task_id) or {}
        relative = str(task.get("generation_spec_path") or "")
        return read_json(self.artifact_path(task_id, relative)) if relative else None

    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 latest_repairable_cdsl(self, task_id: str) -> tuple[str, Path] | None:
        """Return the latest revision CDSL when a quality failure needs repair."""
        task = self.read_task(task_id)
        for revision in reversed((task or {}).get("revisions") or []):
            revision_id = str(revision.get("revision_id") or "")
            if not revision_id or str(revision.get("quality_status") or "") != "needs_repair":
                continue
            path = self.revision_cdsl_path(task_id, revision_id)
            if path is not None:
                return revision_id, path
        return None

    def feature_plan_path(self, task_id: str) -> Path:
        return self.task_dir(task_id) / "feature-plan.json"

    def read_feature_plan(self, task_id: str) -> dict[str, Any] | None:
        safe_task = safe_task_id(task_id)
        plan = read_json(self.feature_plan_path(safe_task))
        if not isinstance(plan, dict):
            return None
        plan_task = str(plan.get("task_id") or "")
        if plan_task and plan_task != safe_task:
            return None
        return plan

    def write_feature_plan(self, task_id: str, plan: dict[str, Any]) -> Path:
        safe_task = safe_task_id(task_id)
        if not isinstance(plan, dict):
            raise ValueError("Feature plan must be an object")
        plan_task = str(plan.get("task_id") or "")
        if plan_task and plan_task != safe_task:
            raise ValueError("Feature plan task_id does not match its task")
        plan["task_id"] = safe_task
        path = self.feature_plan_path(safe_task)
        write_json(path, plan)
        return path

    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