"""doc-gen task: LLM-powered README generation."""

from __future__ import annotations

import difflib

from milco.core.run_context import RunContext
from milco.core.artifacts import write_artifact, ensure_all_artifacts_exist
from milco.tools.fs_tools import scan_repo
from milco.tasks.base import FixResult
from milco.tasks.registry import register_task, TaskBase

MAX_FILE_SAMPLE_CHARS = 2000
MAX_FILES_TO_SAMPLE = 10

SYSTEM_PROMPT = (
    "You are a technical writer. Generate a clear, concise README.md for the given repository. "
    "Include these sections: a title and description, Installation, Usage, and Project Structure. "
    "Use markdown formatting. Be accurate based on the provided context. Do not invent features."
)


def _gather_context(ctx: RunContext) -> str:
    parts: list[str] = []

    repo_map = scan_repo(ctx.repo_root)
    parts.append(f"Repository: {repo_map['repo_root']}")
    parts.append(f"Total files: {repo_map['total_files']}")
    parts.append(f"Python files: {repo_map['python_file_count']}")
    parts.append("\nFile listing:\n" + "\n".join(f"  {f}" for f in repo_map["files"]))

    pyproject = ctx.repo_root / "pyproject.toml"
    if pyproject.exists():
        parts.append(
            f"\n--- pyproject.toml ---\n{pyproject.read_text(encoding='utf-8')}"
        )

    existing_readme = ctx.repo_root / "README.md"
    if existing_readme.exists():
        parts.append(
            f"\n--- Existing README.md ---\n{existing_readme.read_text(encoding='utf-8')}"
        )

    py_files = [
        f for f in repo_map["files"] if f.endswith(".py") and not f.startswith("tests/")
    ]
    for rel_path in py_files[:MAX_FILES_TO_SAMPLE]:
        full = ctx.repo_root / rel_path
        if full.exists():
            content = full.read_text(encoding="utf-8")[:MAX_FILE_SAMPLE_CHARS]
            parts.append(f"\n--- {rel_path} ---\n{content}")

    return "\n".join(parts)


def _make_unified_diff(old_content: str, new_content: str, filename: str) -> str:
    old_lines = old_content.splitlines(keepends=True)
    new_lines = new_content.splitlines(keepends=True)
    diff = difflib.unified_diff(
        old_lines,
        new_lines,
        fromfile=f"a/{filename}",
        tofile=f"b/{filename}",
    )
    return "".join(diff)


@register_task("doc-gen")
class DocGenTask(TaskBase):
    needs_llm = True
    description = "Generate/update README.md from repo context"

    def fix(self, ctx: RunContext) -> FixResult:
        ctx.ensure_run_dir()
        errors: list[str] = []

        context = _gather_context(ctx)
        prompt = f"Generate a README.md for this repository:\n\n{context}"

        try:
            readme_content = ctx.llm.complete(prompt, system=SYSTEM_PROMPT)
        except Exception as exc:
            errors.append(f"LLM call failed: {exc}")
            write_artifact(ctx, "patches.diff", f"# LLM call failed: {exc}\n")
            ensure_all_artifacts_exist(ctx)
            return FixResult(success=False, applied=False, diff_text="", errors=errors)

        readme_path = ctx.repo_root / "README.md"
        old_content = (
            readme_path.read_text(encoding="utf-8") if readme_path.exists() else ""
        )
        diff_text = _make_unified_diff(old_content, readme_content, "README.md")

        if not diff_text.strip():
            diff_text = "# No changes needed – README.md is up to date.\n"

        write_artifact(ctx, "patches.diff", diff_text)

        if ctx.can_apply():
            readme_path.write_text(readme_content, encoding="utf-8")
            applied = True
        else:
            applied = False

        ensure_all_artifacts_exist(ctx)
        return FixResult(
            success=len(errors) == 0,
            applied=applied,
            diff_text=diff_text,
            errors=errors,
        )
