import asyncio
import contextvars
import inspect
import json
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

MIXER_DIR = Path(__file__).parent.parent.resolve()
_ENV_LOADED_ROOTS: set[str] = set()

# --- 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"),
        }

    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:
                            task_hint = str(Path(self.session_folder) / "task.md") if self.session_folder else "task.md"
                            raise ValueError(
                                f"Unknown module '{target}' in {task_hint}. "
                                "Update modules to valid module names."
                            )
                        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:
                            task_hint = str(Path(self.session_folder) / "task.md") if self.session_folder else "task.md"
                            raise ValueError(
                                f"Unknown module '{target}' in {task_hint}. "
                                "Update modules to valid module names."
                            )
                        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) -> str:
        """List all .md files in the session folder as prior artifacts.

        Returns formatted bullet list of existing .md paths,
        or '(none)' if the session folder has no .md files.
        """
        if not self.session_folder:
            return "   (none)"
        wf = Path(self.session_folder)
        md_files = sorted(wf.glob("*.md"))
        if not md_files:
            return "   (none)"
        lines = [f"   - {rel_path(str(p))}" for p in md_files]
        return "\n".join(lines)

    def resolve_questions(self, stage: str) -> str:
        """Read answered, non-implemented questions for a stage.
        Returns formatted text or empty string."""
        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_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)


# --- Task parser ---


@dataclass
class TaskData:
    title: str = ""
    body: str = ""


def parse_task(task_path: str) -> TaskData:
    """Parse a task.md file into a TaskData object.

    Reads clean markdown (no frontmatter). Title comes from the first # heading.
    Modules come from session.json (not from this file).
    """
    text = Path(task_path).read_text()

    title = ""
    title_match = re.search(r"^# (.+)$", text, re.MULTILINE)
    if title_match:
        title = title_match.group(1).strip()

    return TaskData(
        title=title,
        body=text,
    )


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