import asyncio
import contextvars
import inspect
import json
import os
import re
import subprocess
import sys
from contextlib import contextmanager
from dataclasses import dataclass, field
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
    work_folder: str | None = None
    size: str = "M"
    depth: int = 0

    @property
    def paths(self) -> dict[str, str]:
        """Standard work_folder paths: task_path, plan_path, work_path, update_path, upgrade_path, agent_logs_path."""
        if not self.work_folder:
            return {}
        wf = Path(self.work_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"),
            "agent_logs_path": str(wf / "agent_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 doc.md paths from settings.json tree (ancestor-aware)
        - "rules": collects .mixer/modules/<name>/<action>.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.work_folder) / "task.md") if self.work_folder else "task.md"
                            raise ValueError(
                                f"Unknown module '{target}' in {task_hint}. "
                                "Update task frontmatter 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.work_folder) / "task.md") if self.work_folder else "task.md"
                            raise ValueError(
                                f"Unknown module '{target}' in {task_hint}. "
                                "Update task frontmatter modules to valid module names."
                            )
                        names.extend(chain)
                for name in names:
                    rule_path = project_root / ".mixer" / "modules" / name / f"{action}.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 work folder as prior artifacts.

        Returns formatted bullet list of existing .md paths,
        or '(none)' if the work folder has no .md files.
        """
        if not self.work_folder:
            return "   (none)"
        wf = Path(self.work_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_docs(self) -> str:
        """Resolve doc.md paths 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 task.md in the current work_folder."""
        if not self.work_folder:
            return ["all"]
        task_path = Path(self.work_folder) / "task.md"
        if not task_path.is_file():
            return ["all"]
        return parse_task(str(task_path)).modules  # parse_task is defined below


_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(work_folder: str, size: str = "M", depth: int = 0,
                project_root: str | None = None):
    """Set workflow context directly without automatic reset."""
    root = _resolve_project_root(project_root)
    wf = Path(work_folder)
    if not wf.is_absolute():
        wf = root / wf
    _load_project_env(root)
    _ctx_var.set(WorkflowContext(
        project_root=str(root),
        work_folder=str(wf.resolve()),
        size=size,
        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,
                    work_folder: str | None,
                    size: str | 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_work_folder = current.work_folder
    if work_folder is not None:
        wf_path = Path(work_folder)
        if not wf_path.is_absolute():
            wf_path = Path(next_project_root) / wf_path
        next_work_folder = str(wf_path.resolve())

    next_size = size if size is not None else current.size

    if depth is not None:
        next_depth = depth
    elif current.work_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,
        work_folder=next_work_folder,
        size=next_size,
        depth=next_depth,
    )


@contextmanager
def context_scope(*,
                  project_root: str | None = None,
                  work_folder: str | None = None,
                  size: str | 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,
        work_folder=work_folder,
        size=size,
        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:
        raise ValueError(
            f"Invalid modules value in {task_path}: modules cannot be empty. "
            "Use 'all' or a non-empty module list."
        )

    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 == "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 set_active_work_folder(project_root: Path, work_folder: str) -> None:
    """Write the active work folder path to .mixer/active-work-folder state file."""
    state_file = project_root / ".mixer" / "active-work-folder"
    state_file.parent.mkdir(parents=True, exist_ok=True)
    state_file.write_text(work_folder)


def clear_active_work_folder(project_root: Path) -> None:
    """Remove the .mixer/active-work-folder state file if it exists."""
    state_file = project_root / ".mixer" / "active-work-folder"
    if state_file.is_file():
        state_file.unlink()


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)

    work_folder = args.pop("work_folder", None)
    if not work_folder:
        print("Usage: --work_folder=<path>")
        sys.exit(1)
    wf_path = Path(work_folder)
    if not wf_path.is_absolute():
        wf_path = project_root / wf_path
    work_folder = str(wf_path.resolve())

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

            with context_scope(
                project_root=str(project_root),
                work_folder=work_folder,
                size=size,
                depth=0,
                increment_depth=False,
            ):
                result = asyncio.run(run_fn(**args))
    finally:
        clear_active_work_folder(project_root)

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


# --- Task parser ---


@dataclass
class TaskData:
    name: str = ""
    modules: list[str] = field(default_factory=lambda: ["all"])
    assignee: str = ""
    cycle: str = ""
    status: str = ""
    title: str = ""
    body: str = ""


def _parse_frontmatter(text: str) -> tuple[dict, str]:
    """Split text into (frontmatter_dict, body)."""
    lines = text.splitlines(keepends=True)
    if not lines or lines[0].strip() != "---":
        return {}, text

    end_index = None
    for i in range(1, len(lines)):
        if lines[i].strip() == "---":
            end_index = i
            break

    if end_index is None:
        return {}, text

    raw_lines = [line.rstrip("\r\n") for line in lines[1:end_index]]
    body = "".join(lines[end_index + 1:]).lstrip("\n")

    meta = {}
    i = 0
    while i < len(raw_lines):
        line = raw_lines[i]
        stripped = line.strip()
        if not stripped or stripped.startswith("#"):
            i += 1
            continue

        match = re.match(r"^([A-Za-z0-9_-]+)\s*:\s*(.*)$", stripped)
        if not match:
            i += 1
            continue

        key = match.group(1)
        value = match.group(2).strip()

        # Support minimal YAML list style:
        # modules:
        #   - api
        #   - web
        if value == "":
            items = []
            j = i + 1
            while j < len(raw_lines):
                child = raw_lines[j]
                child_stripped = child.strip()
                if not child_stripped:
                    j += 1
                    continue
                if not child.startswith((" ", "\t")):
                    break
                item_match = re.match(r"^\s*-\s+(.*)$", child)
                if not item_match:
                    break
                item = item_match.group(1).strip()
                if (item.startswith('"') and item.endswith('"')) or (
                    item.startswith("'") and item.endswith("'")
                ):
                    item = item[1:-1]
                items.append(item)
                j += 1

            if items:
                meta[key] = items
                i = j
                continue

            meta[key] = ""
            i += 1
            continue

        if (value.startswith('"') and value.endswith('"')) or (
            value.startswith("'") and value.endswith("'")
        ):
            value = value[1:-1]

        meta[key] = value
        i += 1

    return meta, body


def _serialize_frontmatter(meta: dict) -> str:
    """Convert a dict to YAML frontmatter string (with --- delimiters)."""
    def _yaml_scalar(value: str) -> str:
        if value == "":
            return '""'
        if re.fullmatch(r"[A-Za-z0-9._/@+\-]+", value):
            return value
        # JSON quoting is YAML-compatible for simple scalar strings.
        return json.dumps(value)

    lines = ["---"]
    for key, value in meta.items():
        if isinstance(value, list):
            if not value:
                lines.append(f"{key}: []")
                continue
            lines.append(f"{key}:")
            for item in value:
                lines.append(f"  - {_yaml_scalar(str(item))}")
            continue
        lines.append(f"{key}: {_yaml_scalar(str(value))}")
    lines.append("---")
    return "\n".join(lines)


def _parse_legacy(text: str) -> tuple[dict, str]:
    """Fallback parser for old format: # Task: <title>."""
    meta = {}

    title_match = re.search(r"^# Task:\s*(.+)$", text, re.MULTILINE)
    if title_match:
        meta["_title"] = title_match.group(1).strip()

    return meta, text


def parse_task(task_path: str) -> TaskData:
    """Parse a task.md file into a TaskData object."""
    text = Path(task_path).read_text()
    meta, body = _parse_frontmatter(text)

    legacy_title = ""
    if not meta:
        legacy_meta, body = _parse_legacy(text)
        meta = legacy_meta
        legacy_title = meta.pop("_title", "")

    raw_modules = meta.get("modules", "all")
    modules = _normalize_modules(raw_modules, task_path)

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

    return TaskData(
        name=meta.get("name", ""),
        modules=modules,
        assignee=meta.get("assignee", ""),
        cycle=meta.get("cycle", ""),
        status=meta.get("status", ""),
        title=title,
        body=body,
    )


def update_frontmatter(file_path: str, **kwargs) -> None:
    """Update specific keys in the YAML frontmatter of a markdown file."""
    text = Path(file_path).read_text()
    meta, body = _parse_frontmatter(text)

    if not meta and not text.startswith("---"):
        meta = {}
        body = text

    for key, value in kwargs.items():
        if isinstance(value, list):
            meta[key] = [str(v) for v in value]
        elif isinstance(value, bool):
            meta[key] = "true" if value else "false"
        elif isinstance(value, int):
            meta[key] = str(value)
        else:
            meta[key] = str(value)

    new_text = _serialize_frontmatter(meta) + "\n\n" + body
    Path(file_path).write_text(new_text)


parse_frontmatter = _parse_frontmatter
serialize_frontmatter = _serialize_frontmatter


# --- Git helpers ---


def is_git_repo(cwd: str) -> bool:
    """Check if cwd is inside a git repo."""
    result = subprocess.run(["git", "rev-parse", "--is-inside-work-tree"],
                            cwd=cwd, capture_output=True, text=True)
    return result.returncode == 0


def git_checkpoint(cwd: str):
    """Commit all current changes as a checkpoint. Skips if nothing to commit."""
    result = subprocess.run(["git", "status", "--porcelain"],
                            cwd=cwd, capture_output=True, text=True)
    if not result.stdout.strip():
        return
    subprocess.run(["git", "add", "."], cwd=cwd, check=True)
    subprocess.run(["git", "commit", "-m", "checkpoint: before build"],
                   cwd=cwd, check=True)


def git_commit_and_push(cwd: str):
    """Commit build result and push."""
    result = subprocess.run(["git", "status", "--porcelain"],
                            cwd=cwd, capture_output=True, text=True)
    if not result.stdout.strip():
        return
    subprocess.run(["git", "add", "."], cwd=cwd, check=True)
    subprocess.run(["git", "commit", "-m", "build: completed"],
                   cwd=cwd, check=True)
    subprocess.run(["git", "push"], cwd=cwd, check=True)
