from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass

import aiosqlite

from mcp.server.fastmcp import Context, FastMCP

from peon_mcp.db import init_db, row_to_dict


# ---------------------------------------------------------------------------
# Lifespan – keeps a single DB connection for the server's lifetime
# ---------------------------------------------------------------------------


@dataclass
class AppContext:
    db: aiosqlite.Connection


@asynccontextmanager
async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
    db = await init_db()
    try:
        yield AppContext(db=db)
    finally:
        await db.close()


# ---------------------------------------------------------------------------
# Server
# ---------------------------------------------------------------------------

mcp = FastMCP("peon-mcp", lifespan=app_lifespan)


# ---------------------------------------------------------------------------
# Prompts
# ---------------------------------------------------------------------------


@mcp.prompt()
def start_work(project_id: str, repo_path: str) -> str:
    """Bootstrap prompt that tells an agent how to begin working on a project.

    The agent will: ensure the project is registered, call next_task to pick
    up work, then use the execute_task prompt to get its full instructions.

    Args:
        project_id: The repo name (e.g. 'peon-mcp').
        repo_path: Absolute path to the repo on disk.
    """
    return f"""\
You are a Peon — a hardworking Orc laborer from Warcraft. You speak in Peon mannerisms ('Zug zug', 'Job's done!', 'Work work', 'Dabu'). Your work logs should reflect this persona, but your actual code and commits must remain professional and high quality. The persona lives in your summaries and logs, NOT in the code you write for the project.

You are a development agent about to start work on project '{project_id}' \
located at {repo_path}.

Follow these steps in order:

1. Call get_project(project_id='{project_id}'). If the project does not \
exist, call create_project(project_id='{project_id}', path='{repo_path}') \
to register it.

2. **Clean working tree** — Run `git status` in {repo_path}. If there are \
untracked or uncommitted changes, run `git stash --include-untracked` to \
stash them and call log_work to record that you stashed pre-existing changes \
(include the stash ref in the summary).

3. Run `git rev-parse HEAD` in {repo_path} to capture the current commit SHA.

4. Call next_task(project_id='{project_id}', commit_base='<sha>') passing \
the SHA from step 3. This marks the task as in_progress, records the base \
commit, and returns the task details.

5. Use the execute_task prompt with project_id='{project_id}' and the task's \
title and description to get your full working instructions, then follow them.\
"""


@mcp.prompt()
def execute_task(project_id: str, task_title: str, task_description: str) -> str:
    """Instructions for an agent to execute a single task.

    Intended to be combined with the output of next_task so the orchestrator
    can hand an agent both the work item and the rules for doing it.

    Args:
        project_id: The repo name (e.g. 'peon-mcp').
        task_title: Title of the task to execute.
        task_description: Full description of the task.
    """
    return f"""\
You are a development agent working on project '{project_id}'.

## Your current assignment

**{task_title}**
{task_description}

## Rules

1. **Scope** — Only make changes that directly address the task above. \
Do not refactor surrounding code or add unrelated improvements.

2. **Understand first** — Read the relevant source files before making \
changes. Understand the existing patterns and conventions in the codebase.

3. **Testing** — Before marking the task done:
   - Run the existing test suite and confirm it passes.
   - If you added new behaviour, add tests that cover it.
   - If you modified existing behaviour, update affected tests.
   - Never leave the test suite in a broken state.

4. **If something breaks** — If your changes cause test failures, fix them \
before finishing. If you cannot fix them, revert your changes and update \
the task with a description of what went wrong so it can be retried.

5. **Follow-up work** — If you discover work that is out of scope for this \
task, call create_task to capture it rather than doing it now.

6. **Completion** — When the task is done and tests pass, run \
`git rev-parse HEAD` to capture the final commit SHA, then call \
update_task(task_id=<id>, status='done', commit_head='<sha>'). \
If completing this task also finishes a feature, call \
update_feature(feature_id=<id>, status='done').

7. **Commit and push** — Once tests pass, stage your changes, commit with a \
clear message referencing the task, and push to the remote. Do not amend \
existing commits. Do not push if tests are failing.

8. **Log your work** — After completing (or failing) the task, call \
log_work with:
   - task_id: the task you worked on
   - project_id: '{project_id}'
   - summary: what you did and the outcome
   - files_changed: comma-separated list of files you touched
   - test_result: 'pass', 'fail', or 'skip'
Write your summary in the voice of a Warcraft Peon — use phrases like \
'Zug zug', 'Job\\'s done!', 'Work work' etc. Keep the technical details \
accurate and clear, but add personality.
Always log, even if the task failed — the log should capture what went wrong.\
"""


# ---------------------------------------------------------------------------
# Project tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def create_project(
    project_id: str,
    path: str,
    ctx: Context = None,
) -> dict:
    """Register a project so tasks can be scoped to it.

    Args:
        project_id: Unique identifier for the project (typically the repo name, e.g. 'peon-mcp').
        path: Absolute path to the project on disk (e.g. '/Users/me/projects/peon-mcp').
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    await db.execute(
        "INSERT INTO projects (id, path) VALUES (?, ?)",
        (project_id, path),
    )
    await db.commit()
    rows = await db.execute_fetchall(
        "SELECT * FROM projects WHERE id = ?", (project_id,)
    )
    return row_to_dict(rows[0])


@mcp.tool()
async def list_projects(ctx: Context = None) -> list[dict]:
    """List all registered projects."""
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    rows = await db.execute_fetchall(
        "SELECT * FROM projects ORDER BY created_at DESC"
    )
    return [row_to_dict(r) for r in rows]


@mcp.tool()
async def get_project(project_id: str, ctx: Context = None) -> dict | str:
    """Get a project by its ID.

    Args:
        project_id: The project identifier (repo name).
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    rows = await db.execute_fetchall(
        "SELECT * FROM projects WHERE id = ?", (project_id,)
    )
    if not rows:
        return f"No project found with id '{project_id}'"
    return row_to_dict(rows[0])


# ---------------------------------------------------------------------------
# Feature tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def create_feature(
    project_id: str,
    name: str,
    description: str = "",
    status: str = "planned",
    ctx: Context = None,
) -> dict | str:
    """Create a high-level feature for a project.

    The project must already exist. If it does not, call create_project first.

    Args:
        project_id: The project this feature belongs to (repo name).
        name: Short name for the feature (e.g. 'User authentication').
        description: Longer description of the feature's scope and goals.
        status: One of 'planned', 'in_progress', 'done', or 'cancelled'.
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db

    rows = await db.execute_fetchall(
        "SELECT 1 FROM projects WHERE id = ?", (project_id,)
    )
    if not rows:
        return (
            f"Project '{project_id}' does not exist. "
            f"Call create_project with project_id='{project_id}' and the "
            f"repo path first, then retry create_feature."
        )

    cursor = await db.execute(
        "INSERT INTO features (project_id, name, description, status) VALUES (?, ?, ?, ?)",
        (project_id, name, description, status),
    )
    await db.commit()
    rows = await db.execute_fetchall(
        "SELECT * FROM features WHERE id = ?", (cursor.lastrowid,)
    )
    return row_to_dict(rows[0])


@mcp.tool()
async def list_features(
    project_id: str,
    status: str | None = None,
    ctx: Context = None,
) -> list[dict]:
    """List features for a project, optionally filtered by status.

    Args:
        project_id: The project to list features for (repo name).
        status: Filter by status ('planned', 'in_progress', 'done', 'cancelled'). Omit to show all.
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    query = "SELECT * FROM features"
    params: list[str] = [project_id]
    conditions: list[str] = ["project_id = ?"]

    if status is not None:
        conditions.append("status = ?")
        params.append(status)

    query += " WHERE " + " AND ".join(conditions)
    query += " ORDER BY created_at DESC"

    rows = await db.execute_fetchall(query, params)
    return [row_to_dict(r) for r in rows]


@mcp.tool()
async def get_feature(feature_id: int, ctx: Context = None) -> dict | str:
    """Get a single feature by its ID.

    Args:
        feature_id: The numeric ID of the feature.
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    rows = await db.execute_fetchall(
        "SELECT * FROM features WHERE id = ?", (feature_id,)
    )
    if not rows:
        return f"No feature found with id {feature_id}"
    return row_to_dict(rows[0])


@mcp.tool()
async def update_feature(
    feature_id: int,
    name: str | None = None,
    description: str | None = None,
    status: str | None = None,
    ctx: Context = None,
) -> dict | str:
    """Update fields on an existing feature.

    Args:
        feature_id: The numeric ID of the feature to update.
        name: New name (optional).
        description: New description (optional).
        status: New status — 'planned', 'in_progress', 'done', or 'cancelled' (optional).
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db

    fields: list[str] = []
    params: list = []
    if name is not None:
        fields.append("name = ?")
        params.append(name)
    if description is not None:
        fields.append("description = ?")
        params.append(description)
    if status is not None:
        fields.append("status = ?")
        params.append(status)

    if not fields:
        return "No fields to update"

    fields.append("updated_at = CURRENT_TIMESTAMP")
    params.append(feature_id)

    await db.execute(
        f"UPDATE features SET {', '.join(fields)} WHERE id = ?",
        params,
    )
    await db.commit()

    rows = await db.execute_fetchall(
        "SELECT * FROM features WHERE id = ?", (feature_id,)
    )
    if not rows:
        return f"No feature found with id {feature_id}"
    return row_to_dict(rows[0])


@mcp.tool()
async def delete_feature(feature_id: int, ctx: Context = None) -> str:
    """Delete a feature by its ID.

    Args:
        feature_id: The numeric ID of the feature to delete.
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    cursor = await db.execute("DELETE FROM features WHERE id = ?", (feature_id,))
    await db.commit()
    if cursor.rowcount == 0:
        return f"No feature found with id {feature_id}"
    return f"Feature {feature_id} deleted"


# ---------------------------------------------------------------------------
# Task tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def create_task(
    project_id: str,
    title: str,
    description: str = "",
    priority: str = "medium",
    status: str = "todo",
    commit_base: str = "",
    commit_head: str = "",
    ctx: Context = None,
) -> dict | str:
    """Create a new task within a project.

    The project must already exist. If it does not, call create_project first
    to register it, then retry create_task.

    Args:
        project_id: The project this task belongs to (repo name).
        title: Short summary of the task.
        description: Detailed description of the work to be done.
        priority: One of 'low', 'medium', 'high', or 'critical'.
        status: One of 'todo', 'in_progress', 'done', or 'cancelled'.
        commit_base: Git commit SHA when work began (captured at in_progress).
        commit_head: Git commit SHA after work completed (captured at done).
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db

    rows = await db.execute_fetchall(
        "SELECT 1 FROM projects WHERE id = ?", (project_id,)
    )
    if not rows:
        return (
            f"Project '{project_id}' does not exist. "
            f"Call create_project with project_id='{project_id}' and the "
            f"repo path first, then retry create_task."
        )

    cursor = await db.execute(
        "INSERT INTO tasks (project_id, title, description, priority, status, commit_base, commit_head) VALUES (?, ?, ?, ?, ?, ?, ?)",
        (project_id, title, description, priority, status, commit_base, commit_head),
    )
    await db.commit()
    rows = await db.execute_fetchall(
        "SELECT * FROM tasks WHERE id = ?", (cursor.lastrowid,)
    )
    return row_to_dict(rows[0])


@mcp.tool()
async def list_tasks(
    project_id: str,
    priority: str | None = None,
    status: str | None = None,
    ctx: Context = None,
) -> list[dict]:
    """List tasks for a project, optionally filtered by priority and/or status.

    Args:
        project_id: The project to list tasks for (repo name).
        priority: Filter by priority ('low', 'medium', 'high', 'critical'). Omit to show all.
        status: Filter by status ('todo', 'in_progress', 'done', 'cancelled'). Omit to show all.
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    query = "SELECT * FROM tasks"
    params: list[str] = [project_id]
    conditions: list[str] = ["project_id = ?"]

    if priority is not None:
        conditions.append("priority = ?")
        params.append(priority)
    if status is not None:
        conditions.append("status = ?")
        params.append(status)

    query += " WHERE " + " AND ".join(conditions)

    query += (
        " ORDER BY CASE priority"
        " WHEN 'critical' THEN 0"
        " WHEN 'high' THEN 1"
        " WHEN 'medium' THEN 2"
        " WHEN 'low' THEN 3"
        " END, created_at DESC"
    )

    rows = await db.execute_fetchall(query, params)
    return [row_to_dict(r) for r in rows]


@mcp.tool()
async def get_task(task_id: int, ctx: Context = None) -> dict | str:
    """Get a single task by its ID.

    Args:
        task_id: The numeric ID of the task.
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    rows = await db.execute_fetchall("SELECT * FROM tasks WHERE id = ?", (task_id,))
    if not rows:
        return f"No task found with id {task_id}"
    return row_to_dict(rows[0])


@mcp.tool()
async def next_task(project_id: str, commit_base: str = "", ctx: Context = None) -> dict | str:
    """Pick up the next task to work on for a project.

    Returns the highest-priority task with status 'todo' (critical > high >
    medium > low, oldest first within the same priority) and automatically
    marks it as 'in_progress'.

    Args:
        project_id: The project to pick the next task from (repo name).
        commit_base: Git commit SHA at the time work begins (from git rev-parse HEAD).
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    rows = await db.execute_fetchall(
        "SELECT * FROM tasks"
        " WHERE project_id = ? AND status = 'todo'"
        " ORDER BY CASE priority"
        "   WHEN 'critical' THEN 0"
        "   WHEN 'high' THEN 1"
        "   WHEN 'medium' THEN 2"
        "   WHEN 'low' THEN 3"
        " END, created_at ASC"
        " LIMIT 1",
        (project_id,),
    )
    if not rows:
        return f"No remaining todo tasks for project '{project_id}'"

    task = row_to_dict(rows[0])
    await db.execute(
        "UPDATE tasks SET status = 'in_progress', commit_base = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?",
        (commit_base, task["id"]),
    )
    await db.commit()
    task["status"] = "in_progress"
    task["commit_base"] = commit_base
    return task


@mcp.tool()
async def update_task(
    task_id: int,
    title: str | None = None,
    description: str | None = None,
    priority: str | None = None,
    status: str | None = None,
    commit_base: str | None = None,
    commit_head: str | None = None,
    ctx: Context = None,
) -> dict | str:
    """Update fields on an existing task.

    Args:
        task_id: The numeric ID of the task to update.
        title: New title (optional).
        description: New description (optional).
        priority: New priority — 'low', 'medium', 'high', or 'critical' (optional).
        status: New status — 'todo', 'in_progress', 'done', or 'cancelled' (optional).
        commit_base: Git commit SHA when work began (optional).
        commit_head: Git commit SHA after work completed (optional).
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db

    fields: list[str] = []
    params: list = []
    if title is not None:
        fields.append("title = ?")
        params.append(title)
    if description is not None:
        fields.append("description = ?")
        params.append(description)
    if priority is not None:
        fields.append("priority = ?")
        params.append(priority)
    if status is not None:
        fields.append("status = ?")
        params.append(status)
    if commit_base is not None:
        fields.append("commit_base = ?")
        params.append(commit_base)
    if commit_head is not None:
        fields.append("commit_head = ?")
        params.append(commit_head)

    if not fields:
        return "No fields to update"

    fields.append("updated_at = CURRENT_TIMESTAMP")
    params.append(task_id)

    await db.execute(
        f"UPDATE tasks SET {', '.join(fields)} WHERE id = ?",
        params,
    )
    await db.commit()

    rows = await db.execute_fetchall("SELECT * FROM tasks WHERE id = ?", (task_id,))
    if not rows:
        return f"No task found with id {task_id}"
    return row_to_dict(rows[0])


@mcp.tool()
async def delete_task(task_id: int, ctx: Context = None) -> str:
    """Delete a task by its ID.

    Args:
        task_id: The numeric ID of the task to delete.
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    cursor = await db.execute("DELETE FROM tasks WHERE id = ?", (task_id,))
    await db.commit()
    if cursor.rowcount == 0:
        return f"No task found with id {task_id}"
    return f"Task {task_id} deleted"


# ---------------------------------------------------------------------------
# Work log tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def log_work(
    task_id: int,
    project_id: str,
    summary: str,
    files_changed: str = "",
    test_result: str = "",
    ctx: Context = None,
) -> dict:
    """Log the results of work done on a task.

    Call this after completing (or attempting) a task to create an auditable
    record of what was done.

    Args:
        task_id: The numeric ID of the task this work relates to.
        project_id: The project identifier (repo name).
        summary: A concise description of what was done, what changed, and the outcome.
        files_changed: Comma-separated list of files that were added, modified, or deleted.
        test_result: Overall test suite result — 'pass', 'fail', or 'skip'. Leave empty if tests were not run.
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    cursor = await db.execute(
        "INSERT INTO work_logs (task_id, project_id, summary, files_changed, test_result) VALUES (?, ?, ?, ?, ?)",
        (task_id, project_id, summary, files_changed, test_result),
    )
    await db.commit()
    rows = await db.execute_fetchall(
        "SELECT * FROM work_logs WHERE id = ?", (cursor.lastrowid,)
    )
    return row_to_dict(rows[0])


@mcp.tool()
async def list_work_logs(
    task_id: int | None = None,
    project_id: str | None = None,
    ctx: Context = None,
) -> list[dict]:
    """List work logs, optionally filtered by task or project.

    Args:
        task_id: Filter to logs for a specific task (optional).
        project_id: Filter to logs for a specific project (optional).
    """
    db: aiosqlite.Connection = ctx.request_context.lifespan_context.db
    query = "SELECT * FROM work_logs"
    params: list = []
    conditions: list[str] = []

    if task_id is not None:
        conditions.append("task_id = ?")
        params.append(task_id)
    if project_id is not None:
        conditions.append("project_id = ?")
        params.append(project_id)

    if conditions:
        query += " WHERE " + " AND ".join(conditions)

    query += " ORDER BY created_at DESC"

    rows = await db.execute_fetchall(query, params)
    return [row_to_dict(r) for r in rows]


def main():
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
