import asyncio
import contextvars
import inspect
import json
import logging
import os
import re
import sys
from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from dotenv import load_dotenv

log = logging.getLogger(__name__)

MIXER_DIR = Path(__file__).parent.parent.resolve()
_ENV_LOADED_ROOTS: set[str] = set()
_FAMILY_SKIP_NAMES = {"logs", "branches"}

# --- Family context helpers ---


def _is_session_folder(path: Path) -> bool:
    """Return True if path is a directory that should be treated as a session folder."""
    if not path.is_dir():
        return False
    name = path.name
    if name in _FAMILY_SKIP_NAMES:
        return False
    if name.startswith('.') or name.startswith('_'):
        return False
    return True


def _collect_md_files(folder: Path) -> list[Path]:
    """Collect *.md files from a single folder (non-recursive), sorted."""
    return sorted(folder.glob("*.md"))


def _collect_descendants_md(folder: Path) -> list[Path]:
    """Recursively collect *.md files from all descendant session subfolders."""
    results: list[Path] = []
    for child in sorted(folder.iterdir()):
        if not _is_session_folder(child):
            continue
        results.extend(_collect_md_files(child))
        results.extend(_collect_descendants_md(child))
    return results


# --- Workflow context (set by orchestrators, read by agents and call_agent) ---


@dataclass(frozen=True)
class WorkflowContext:
    """Immutable snapshot of workflow runtime state.

    Pushed by orchestrators via context_scope(), read by agents via get_context().
    Backed by a contextvar for proper asyncio/task isolation.
    """
    project_root: str | None = None
    session_folder: str | None = None
    lite: bool = False
    depth: int = 0

    @property
    def paths(self) -> dict[str, str]:
        """Standard session_folder paths: task_path, plan_path, work_path, update_path, upgrade_path, report_path, logs_path."""
        if not self.session_folder:
            return {}
        wf = Path(self.session_folder)
        return {
            "task_path": str(wf / "task.md"),
            "plan_path": str(wf / "plan.md"),
            "work_path": str(wf / "work.md"),
            "update_path": str(wf / "update.md"),
            "upgrade_path": str(wf / "upgrade.md"),
            "report_path": str(wf / "report.md"),
            "logs_path": str(wf / "logs"),
        }

    @property
    def _sessions_root(self) -> Path | None:
        """Return the .mixer/sessions/ directory, or None if project_root is unset."""
        if not self.project_root:
            return None
        return Path(self.project_root) / ".mixer" / "sessions"

    def resolve_context(self, context: list[str], action: str | None = None) -> str:
        """Resolve CONTEXT list into formatted prompt text for {context} placeholder.

        Reads modules from task.md automatically. Resolution follows modules:
        - "doc": collects module doc paths (_name.md) from settings.json tree (ancestor-aware)
        - "rules": collects .mixer/rules/<action>/<name>.md paths (ancestor-aware)

        Returns formatted bullet list or "" if context is empty.
        """
        if not context:
            return ""
        modules = self._get_modules()
        settings = load_settings()
        tree = settings["modules"]
        project_root = Path(self.project_root) if self.project_root else Path.cwd()
        lines = []
        seen = set()

        for ctx in context:
            paths: list[str] = []

            if ctx == "doc":
                if "all" in modules:
                    paths = _collect_all(tree, "doc")
                else:
                    for target in modules:
                        chain = _find_in_tree(tree, target)
                        if chain is None:
                            log.warning("Unknown module '%s' — skipping doc resolution", target)
                            continue
                        for node in chain:
                            if "doc" in node:
                                paths.append(node["doc"])

            elif ctx == "rules":
                if not action:
                    continue
                if "all" in modules:
                    names = _collect_all_names(tree)
                else:
                    names = []
                    for target in modules:
                        chain = _find_in_tree_with_names(tree, target)
                        if chain is None:
                            log.warning("Unknown module '%s' — skipping rules resolution", target)
                            continue
                        names.extend(chain)
                for name in names:
                    rule_path = project_root / ".mixer" / "rules" / action / f"{name}.md"
                    if rule_path.is_file():
                        paths.append(str(rule_path))

            for p in paths:
                if p not in seen:
                    seen.add(p)
                    lines.append(f"   - {rel_path(p)}")

        return "\n".join(lines)

    def resolve_artifacts(self, exclude_stage: str | None = None) -> str:
        """List .md files in the session folder, excluding the current stage artifact."""
        if not self.session_folder:
            return ""
        wf = Path(self.session_folder)
        exclude = f"{exclude_stage}.md" if exclude_stage else None
        md_files = sorted(p for p in wf.glob("*.md") if p.name != exclude)
        if not md_files:
            return ""
        lines = [f"   - {rel_path(str(p))}" for p in md_files]
        return "\n".join(lines)

    def resolve_family(self) -> str:
        """Collect .md files from parent, sibling, and descendant session folders.

        Rules:
        - Parents: walk upward from session_folder until reaching .mixer/sessions/.
        - Siblings: other folders at the same level (no recursion into their children).
        - Descendants: recurse fully into children/grandchildren/etc.
        - Skip: logs/, branches/, and any folder starting with '.' or '_'.

        Returns formatted sectioned string, or "" if no family artifacts exist.
        """
        if not self.session_folder or not self._sessions_root:
            return ""

        session = Path(self.session_folder).resolve()
        sessions_root = self._sessions_root.resolve()

        try:
            session.relative_to(sessions_root)
        except ValueError:
            return ""

        parent_lines: list[str] = []
        sibling_lines: list[str] = []
        descendant_lines: list[str] = []

        # --- Parents ---
        current = session.parent
        while current != sessions_root and current.is_relative_to(sessions_root):
            md_files = _collect_md_files(current)
            if md_files:
                folder_label = rel_path(str(current))
                parent_lines.append(f"   [{folder_label}]")
                for p in md_files:
                    parent_lines.append(f"   - {rel_path(str(p))}")
            current = current.parent

        # --- Siblings ---
        parent_dir = session.parent
        if parent_dir.is_dir() and parent_dir != sessions_root.parent:
            for sibling in sorted(parent_dir.iterdir()):
                if sibling.resolve() == session:
                    continue
                if not _is_session_folder(sibling):
                    continue
                md_files = _collect_md_files(sibling)
                if md_files:
                    folder_label = rel_path(str(sibling))
                    sibling_lines.append(f"   [{folder_label}]")
                    for p in md_files:
                        sibling_lines.append(f"   - {rel_path(str(p))}")

        # --- Descendants ---
        descendant_md = _collect_descendants_md(session)
        if descendant_md:
            for p in descendant_md:
                descendant_lines.append(f"   - {rel_path(str(p))}")

        # Assemble output
        sections: list[str] = []
        if parent_lines:
            sections.append("Parent sessions:\n" + "\n".join(parent_lines))
        if sibling_lines:
            sections.append("Sibling sessions:\n" + "\n".join(sibling_lines))
        if descendant_lines:
            sections.append("Child/descendant sessions:\n" + "\n".join(descendant_lines))

        if not sections:
            return ""
        return "\n\n".join(sections)

    def resolve_questions(self, stage: str) -> str:
        """Read answered, non-implemented questions for a stage."""
        if not self.session_folder:
            return ""
        path = Path(self.session_folder) / f"{stage}_questions.json"
        if not path.is_file():
            return ""
        try:
            data = json.loads(path.read_text())
        except (json.JSONDecodeError, OSError):
            return ""
        answered = [q for q in data.get("questions", [])
                    if q.get("answer") and not q.get("implemented")]
        if not answered:
            return ""
        lines = []
        for q in answered:
            lines.append(f"Q: {q['text']}")
            lines.append(f"A: {q['answer']}")
            lines.append("")
        return "\n".join(lines)

    def resolve_current_artifact(self, stage: str) -> str:
        """Return the current artifact path for a stage if it exists."""
        if not self.session_folder:
            return ""
        path = Path(self.session_folder) / f"{stage}.md"
        if not path.is_file():
            return ""
        return f"   - {rel_path(str(path))}"

    def resolve_role(self, stage: str, role_type: str) -> str:
        """Extract a role section from the stage prompts file.

        role_type is 'artifact-builder' or 'question-builder'. Returns the text
        inside the matching XML tag, e.g. <artifact-builder-role>...</artifact-builder-role>.
        """
        prompts_path = MIXER_DIR / "workflows" / stage / f"{stage}-prompts.md"
        if not prompts_path.is_file():
            return ""
        text = prompts_path.read_text()
        open_tag = f"<{role_type}-role>"
        close_tag = f"</{role_type}-role>"
        start = text.find(open_tag)
        if start == -1:
            return ""
        end = text.find(close_tag, start)
        if end == -1:
            return ""
        return text[start + len(open_tag):end].strip()

    def resolve_template(self, stage: str) -> str:
        """Resolve the template path for a given stage into a formatted bullet."""
        template_path = MIXER_DIR / "workflows" / stage / f"{stage}-template.md"
        return f"   - {rel_path(str(template_path))}"

    def resolve_docs(self) -> str:
        """Resolve module doc paths (_name.md) into formatted prompt text."""
        return self.resolve_context(["doc"])

    def resolve_rules(self, action: str) -> str:
        """Resolve rules paths into formatted prompt text."""
        return self.resolve_context(["rules"], action=action)

    def _get_modules(self) -> list[str]:
        """Read modules from session.json."""
        if not self.session_folder:
            return []
        session_data = read_session_json(self.session_folder)
        modules = session_data.get("modules")
        if modules:
            return _normalize_modules(modules, self.session_folder + "/session.json")
        return []


_ctx_var: contextvars.ContextVar[WorkflowContext] = contextvars.ContextVar(
    "workflow_context", default=WorkflowContext()
)


def _resolve_project_root(project_root: str | Path | None = None) -> Path:
    if project_root is not None:
        root = Path(project_root)
        if not root.is_absolute():
            root = (Path.cwd() / root).resolve()
        return root.resolve()
    current = get_context().project_root
    if current:
        return Path(current).resolve()
    return Path.cwd().resolve()


def _load_project_env(project_root: Path) -> None:
    key = str(project_root.resolve())
    if key in _ENV_LOADED_ROOTS:
        return
    load_dotenv(project_root / ".env")
    _ENV_LOADED_ROOTS.add(key)


def get_project_root() -> Path:
    return _resolve_project_root()


def set_context(session_folder: str, lite: bool = False, depth: int = 0,
                project_root: str | None = None):
    """Set workflow context directly without automatic reset."""
    root = _resolve_project_root(project_root)
    wf = Path(session_folder)
    if not wf.is_absolute():
        wf = root / wf
    _load_project_env(root)
    _ctx_var.set(WorkflowContext(
        project_root=str(root),
        session_folder=str(wf.resolve()),
        lite=lite,
        depth=depth,
    ))


def get_context() -> WorkflowContext:
    """Get the current workflow context. Called by agents and call_agent."""
    return _ctx_var.get()


def _derive_context(*,
                    project_root: str | None,
                    session_folder: str | None,
                    lite: bool | None,
                    depth: int | None,
                    increment_depth: bool) -> WorkflowContext:
    """Build the next context from the current context and optional overrides."""
    current = get_context()

    next_project_root = current.project_root
    if project_root is not None:
        root_path = Path(project_root)
        if not root_path.is_absolute():
            root_path = (Path.cwd() / root_path).resolve()
        next_project_root = str(root_path)
    if not next_project_root:
        next_project_root = str(Path.cwd().resolve())

    next_session_folder = current.session_folder
    if session_folder is not None:
        wf_path = Path(session_folder)
        if not wf_path.is_absolute():
            wf_path = Path(next_project_root) / wf_path
        next_session_folder = str(wf_path.resolve())

    next_lite = lite if lite is not None else current.lite

    if depth is not None:
        next_depth = depth
    elif current.session_folder is not None and increment_depth:
        next_depth = current.depth + 1
    else:
        next_depth = current.depth

    _load_project_env(Path(next_project_root))

    return WorkflowContext(
        project_root=next_project_root,
        session_folder=next_session_folder,
        lite=next_lite,
        depth=next_depth,
    )


@contextmanager
def context_scope(*,
                  project_root: str | None = None,
                  session_folder: str | None = None,
                  lite: bool | None = None,
                  depth: int | None = None,
                  increment_depth: bool = True):
    """Temporarily push a workflow context, then restore the previous one."""
    next_ctx = _derive_context(
        project_root=project_root,
        session_folder=session_folder,
        lite=lite,
        depth=depth,
        increment_depth=increment_depth,
    )
    token = _ctx_var.set(next_ctx)
    try:
        yield next_ctx
    finally:
        _ctx_var.reset(token)


def load_settings() -> dict:
    """Load settings from the active project root."""
    project_root = get_project_root()
    env_override = os.environ.get("MIXERSYSTEM_SETTINGS_PATH")

    candidates = [project_root / ".mixer" / "settings.json"]
    if env_override:
        override_path = Path(env_override)
        if not override_path.is_absolute():
            override_path = project_root / override_path
        candidates.insert(0, override_path)

    seen = set()
    ordered = []
    for path in candidates:
        key = str(path.resolve())
        if key in seen:
            continue
        seen.add(key)
        ordered.append(path)

    for settings_path in ordered:
        if settings_path.is_file():
            return json.loads(settings_path.read_text())

    tried = ", ".join(str(p) for p in ordered)
    raise FileNotFoundError(f"Settings not found. Tried: {tried}")


def _find_in_tree(modules: dict, target: str) -> list[dict] | None:
    """Return ancestor chain [root_node, ..., target_node] or None if not found."""
    for name, node in modules.items():
        if name == target:
            return [node]
        result = _find_in_tree(node.get("children", {}), target)
        if result is not None:
            return [node] + result
    return None


def _collect_all(modules: dict, key: str) -> list[str]:
    """Collect all values for key across the entire tree (DFS)."""
    result = []
    for node in modules.values():
        if key in node:
            result.append(node[key])
        result.extend(_collect_all(node.get("children", {}), key))
    return result


def _find_in_tree_with_names(modules: dict, target: str) -> list[str] | None:
    """Return ancestor name chain [root_name, ..., target_name] or None if not found."""
    for name, node in modules.items():
        if name == target:
            return [name]
        result = _find_in_tree_with_names(node.get("children", {}), target)
        if result is not None:
            return [name] + result
    return None


def _collect_all_names(modules: dict) -> list[str]:
    """Collect all module names across the entire tree (DFS)."""
    result = []
    for name, node in modules.items():
        result.append(name)
        result.extend(_collect_all_names(node.get("children", {})))
    return result


def _dedupe_keep_order(items: list[str]) -> list[str]:
    seen = set()
    out = []
    for item in items:
        if item in seen:
            continue
        seen.add(item)
        out.append(item)
    return out


def _normalize_modules(raw_modules, task_path: str) -> list[str]:
    """Normalize and validate task modules with strict semantics."""
    if isinstance(raw_modules, str):
        stripped = raw_modules.strip()
        parsed_list = None
        if stripped.startswith("[") and stripped.endswith("]"):
            try:
                parsed = json.loads(stripped)
                if isinstance(parsed, list):
                    parsed_list = parsed
            except json.JSONDecodeError:
                parsed_list = None

        if parsed_list is not None:
            modules = [str(m).strip() for m in parsed_list if str(m).strip()]
        else:
            modules = [m.strip() for m in stripped.split(",") if m.strip()]
    elif isinstance(raw_modules, list):
        modules = [str(m).strip() for m in raw_modules if str(m).strip()]
    else:
        raise ValueError(
            f"Invalid modules value in {task_path}: expected string or list, "
            f"got {type(raw_modules).__name__}."
        )

    if not modules:
        return []

    modules = _dedupe_keep_order(modules)
    if "all" in modules:
        if len(modules) > 1:
            raise ValueError(
                f"Invalid modules value in {task_path}: cannot mix 'all' with "
                "specific module names."
            )
        return ["all"]

    return modules


def rel_path(absolute_path: str) -> str:
    """Convert an absolute path to a relative path from the project root.
    Returns a ./ prefixed string. If already relative, returns as-is."""
    p = Path(absolute_path)
    if not p.is_absolute():
        return str(absolute_path)
    try:
        return "./" + str(p.resolve().relative_to(get_project_root()))
    except ValueError:
        return str(absolute_path)


RETRY_FORMAT_PROMPT = """\
Your response was missing the required output format. You MUST end your response with:

[RESULT]
key: value
[NOTES]
Your notes here.

Respond again with the same answer, but include the [RESULT] and [NOTES] blocks at the end."""


@dataclass
class AgentResult:
    """Parsed agent output. Fields from [RESULT] block accessible as attributes."""
    raw: str = ""
    notes: str = ""
    fields: dict = field(default_factory=dict)

    def get(self, key: str, default: str = "") -> str:
        return self.fields.get(key.lower(), default)

    @property
    def status(self) -> str:
        return self.fields.get("status", "").upper()

    @property
    def passed(self) -> bool:
        return self.status in ("APPROVED", "PASS")


def parse_result(text: str) -> AgentResult:
    """Parse standard agent output format."""
    result = AgentResult(raw=text)

    result_marker = text.find("[RESULT]")
    if result_marker == -1:
        return result

    after_result = text[result_marker + len("[RESULT]"):]

    notes_marker = after_result.find("[NOTES]")
    if notes_marker != -1:
        result_section = after_result[:notes_marker]
        result.notes = after_result[notes_marker + len("[NOTES]"):].strip()
    else:
        result_section = after_result

    for line in result_section.strip().splitlines():
        line = line.strip()
        if not line:
            continue
        if ":" in line:
            key, value = line.split(":", 1)
            result.fields[key.strip().lower()] = value.strip()

    return result


def _coerce_arg_value(key: str, value: str):
    """Coerce CLI string value into conservative typed values."""
    lower = value.lower()

    if lower == "true":
        return True
    if lower == "false":
        return False
    if lower in ("none", "null"):
        return None

    if value.startswith("{") or value.startswith("["):
        try:
            parsed = json.loads(value)
            if isinstance(parsed, (dict, list)):
                return parsed
        except json.JSONDecodeError:
            pass

    if key in {"branch_count", "max_revisions"} and re.fullmatch(r"-?\d+", value):
        return int(value)

    if (key.endswith("_count") or key.endswith("_iterations")) and re.fullmatch(r"-?\d+", value):
        return int(value)

    return value


def parse_args(argv: list[str]) -> dict:
    """Parse --key=value arguments from argv."""
    args = {}
    for arg in argv:
        if arg.startswith("--") and "=" in arg:
            key, value = arg[2:].split("=", 1)
            args[key] = _coerce_arg_value(key, value)
        else:
            print(f"Invalid argument: {arg}  (expected --key=value)")
            sys.exit(1)
    return args


def run_cli(run_fn):
    """Parse CLI args, invoke run_fn, and print structured results."""
    args = parse_args(sys.argv[1:])
    # depth is internal-only; never accept it from public CLI input.
    args.pop("depth", None)
    project_root_arg = args.pop("project_root", None)
    project_root = _resolve_project_root(project_root_arg)
    _load_project_env(project_root)

    session_folder = args.pop("session_folder", None)
    if not session_folder:
        print("Usage: --session_folder=<path>")

        sys.exit(1)
    wf_path = Path(session_folder)
    if not wf_path.is_absolute():
        wf_path = project_root / wf_path
    session_folder = str(wf_path.resolve())

    try:
        sig = inspect.signature(run_fn)
        if "session_folder" in sig.parameters:
            with context_scope(
                project_root=str(project_root),
                increment_depth=False,
            ):
                result = asyncio.run(run_fn(session_folder, **args))
        else:
            lite = args.pop("lite", False)

            with context_scope(
                project_root=str(project_root),
                session_folder=session_folder,
                lite=lite,
                depth=0,
                increment_depth=False,
            ):
                result = asyncio.run(run_fn(**args))
    finally:
        pass

    if result is None:
        return
    if isinstance(result, (dict, list)):
        print(json.dumps(result, indent=2))
    else:
        print(result)



def _parse_frontmatter_block(text: str) -> tuple[dict, str]:
    """Parse a simple YAML-like frontmatter block from markdown text."""
    if not text.startswith("---\n"):
        return {}, text

    closing = text.find("\n---\n", 4)
    if closing < 0:
        return {}, text

    raw_frontmatter = text[4:closing]
    body = text[closing + 5:]

    data: dict = {}
    current_list_key: str | None = None
    for line in raw_frontmatter.splitlines():
        stripped = line.strip()
        if not stripped:
            continue
        if line.startswith("  - ") and current_list_key:
            data.setdefault(current_list_key, []).append(stripped[2:].strip())
            continue
        if ":" not in stripped:
            current_list_key = None
            continue
        key, value = stripped.split(":", 1)
        key = key.strip()
        value = value.strip()
        if not value:
            data[key] = []
            current_list_key = key
            continue
        current_list_key = None
        if value.lower() == "true":
            data[key] = True
        elif value.lower() == "false":
            data[key] = False
        elif re.fullmatch(r"-?\d+", value):
            data[key] = int(value)
        else:
            data[key] = value
    return data, body


def _render_frontmatter_block(data: dict) -> str:
    lines: list[str] = []
    for key, value in data.items():
        if isinstance(value, list):
            lines.append(f"{key}:")
            for item in value:
                lines.append(f"  - {item}")
            continue
        if isinstance(value, bool):
            rendered = "true" if value else "false"
        else:
            rendered = str(value)
        lines.append(f"{key}: {rendered}")
    return "\n".join(lines)


def update_frontmatter(task_path: str, **kwargs) -> None:
    """Compatibility helper to upsert frontmatter keys in markdown files."""
    path = Path(task_path)
    raw = path.read_text() if path.is_file() else ""
    frontmatter, body = _parse_frontmatter_block(raw)
    frontmatter.update(kwargs)
    rendered = _render_frontmatter_block(frontmatter)
    body = body.lstrip("\n")
    path.write_text(f"---\n{rendered}\n---\n\n{body}")


# --- session.json helpers ---


def init_session(session_folder: str) -> dict:
    """Initialize session.json with default modules if it doesn't exist yet.

    Seeds modules from settings.json default_modules and sets modules_locked=false.
    If session.json already exists, returns it as-is.
    """
    p = Path(session_folder) / "session.json"
    if p.is_file():
        return json.loads(p.read_text())
    settings = load_settings()
    defaults = settings.get("default_modules", [])
    data = {"modules": defaults, "modules_locked": False}
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(data, indent=2) + "\n")
    return data


def read_session_json(session_folder: str) -> dict:
    """Read session.json from a session folder. Returns {} if missing."""
    p = Path(session_folder) / "session.json"
    if not p.is_file():
        return {}
    return json.loads(p.read_text())


def update_session_json(session_folder: str, **kwargs) -> dict:
    """Merge top-level keys into session.json. Creates file if missing."""
    p = Path(session_folder) / "session.json"
    p.parent.mkdir(parents=True, exist_ok=True)
    data = json.loads(p.read_text()) if p.is_file() else {}
    for key, value in kwargs.items():
        if isinstance(value, dict) and isinstance(data.get(key), dict):
            data[key] = {**data[key], **value}  # shallow merge for nested dicts
        else:
            data[key] = value
    p.write_text(json.dumps(data, indent=2) + "\n")
    return data


def update_session_run(session_folder: str, stage: str, **kwargs) -> None:
    """Append a run entry to runs.<stage> in session.json."""
    data = read_session_json(session_folder)
    runs = data.setdefault("runs", {})
    entries = runs.setdefault(stage, [])
    if not isinstance(entries, list):
        entries = []
        runs[stage] = entries
    entry = {"at": datetime.now(timezone.utc).isoformat(), **kwargs}
    entries.append(entry)
    Path(session_folder, "session.json").parent.mkdir(parents=True, exist_ok=True)
    Path(session_folder, "session.json").write_text(json.dumps(data, indent=2) + "\n")
