import asyncio
import re
import tempfile
from contextlib import asynccontextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

import regex
import yaml
from agents import RunContextWrapper, function_tool

from config import WORKSPACE
from core.runtime.context import AgentRuntimeContext
from schema.common.tool_results import ToolResultSchema, ToolResultStatusSchema, ToolResultTypeSchema


_KNOWLEDGE_NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$")
KNOWLEDGES_DIR_NAME = "knowledges"
KNOWLEDGE_EXTENSION = ".md"
_MAX_DESCRIPTION_LENGTH = 512
_MAX_CONTENT_LENGTH = 200_000
_MAX_FIND_KEYWORD_LENGTH = 200
_DEFAULT_FIND_MAX_MATCHES = 50
_MAX_FIND_MAX_MATCHES = 200
_DEFAULT_FIND_CONTEXT_LINES = 0
_MAX_FIND_CONTEXT_LINES = 5
_REGEX_SEARCH_TIMEOUT_SECONDS = 0.05
_knowledge_locks: dict[str, asyncio.Lock] = {}
_knowledge_generation = 0


@dataclass(slots=True)
class KnowledgeDocument:
    name: str
    description: str
    body: str

    @classmethod
    def parse(cls, text: str) -> "KnowledgeDocument":
        front_matter, body = _split_front_matter(text)
        metadata = yaml.safe_load(front_matter) if front_matter else None
        if not isinstance(metadata, dict):
            raise ValueError("Knowledge Front Matter must be a YAML mapping.")
        if set(metadata) != {"name", "description"}:
            raise ValueError("Knowledge Front Matter must contain only name and description.")
        name = metadata.get("name")
        description = metadata.get("description")
        if not isinstance(name, str) or not _KNOWLEDGE_NAME_PATTERN.fullmatch(name):
            raise ValueError("Knowledge Front Matter name is invalid.")
        if not isinstance(description, str) or not description.strip():
            raise ValueError("Knowledge Front Matter description is required.")
        return cls(name=name.strip(), description=description.strip(), body=body)

    def render(self) -> str:
        metadata = yaml.safe_dump(
            {"name": self.name, "description": self.description},
            allow_unicode=True,
            default_flow_style=False,
            sort_keys=False,
        ).strip()
        body = self.body.rstrip("\n")
        return f"---\n{metadata}\n---\n\n{body}\n"

    def body_lines(self) -> list[str]:
        return self.body.split("\n") if self.body else []

    def slice_body(self, start_line: int | None, end_line: int | None) -> tuple[int, list[str]]:
        lines = self.body_lines()
        start, end = _normalize_line_range(start_line, end_line, len(lines), allow_default=True)
        if start is None or end is None:
            return 1, lines
        return start, lines[start - 1:end]

    def patch_body(self, start_line: int, end_line: int, content: str) -> None:
        lines = self.body_lines()
        start, end = _normalize_line_range(start_line, end_line, len(lines), allow_default=False)
        replacement = content.splitlines()
        lines[start - 1:end] = replacement
        self.body = "\n".join(lines)


def current_knowledge_generation() -> int:
    return _knowledge_generation


def _bump_knowledge_generation() -> None:
    global _knowledge_generation
    _knowledge_generation += 1


@function_tool
async def create_knowledge(
    ctx: RunContextWrapper[AgentRuntimeContext],
    name: str,
    description: str,
    content: str,
) -> str:
    """Create a knowledge file for the current agent.

    Args:
        name: str knowledge file name without .md under the current agent's knowledges directory.
        description: str brief knowledge metadata description.
        content: str markdown body content. YAML Front Matter is generated by the tool.

    Returns:
        JSON status describing the creation result.
    """
    try:
        agent_code = _valid_agent_code(ctx)
        knowledge_name = _valid_knowledge_name(name)
        body = _valid_content(content)
        doc = KnowledgeDocument(
            name=knowledge_name,
            description=_valid_description(description),
            body=body,
        )
        knowledge_path = _safe_knowledge_path(agent_code, knowledge_name)
        async with _locked_knowledge_paths(knowledge_path):
            if knowledge_path.exists():
                raise ValueError(f"Knowledge already exists for agent {agent_code}: {knowledge_name}")
            knowledge_path.parent.mkdir(parents=True, exist_ok=True)
            _write_text_atomic(knowledge_path, doc.render())
        _bump_knowledge_generation()
    except Exception as exc:
        return _tool_error(ToolResultTypeSchema.KNOWLEDGE_MUTATION, str(exc) or "Knowledge creation failed.")

    return _tool_success(
        ToolResultTypeSchema.KNOWLEDGE_MUTATION,
        f"Knowledge created for agent {agent_code}: {knowledge_name}",
    )


@function_tool
async def load_knowledge(
    ctx: RunContextWrapper[AgentRuntimeContext],
    name: str,
    start_line: int | None = None,
    end_line: int | None = None,
) -> str:
    """Load the body of a named knowledge file for the current agent.

    Args:
        name: str knowledge file name without .md under the current agent's knowledges directory.
        start_line: int | None optional 1-based body line to start reading from.
        end_line: int | None optional 1-based body line to stop reading at, inclusive.

    Returns:
        JSON status with line-numbered markdown body content without YAML Front Matter.
    """
    try:
        agent_code = _valid_agent_code(ctx)
        knowledge_name = _valid_knowledge_name(name)
        knowledge_path = _existing_knowledge_path(agent_code, knowledge_name)
        doc = KnowledgeDocument.parse(knowledge_path.read_text(encoding="utf-8"))
        first_line, lines = doc.slice_body(start_line, end_line)
    except Exception as exc:
        return _tool_error(ToolResultTypeSchema.KNOWLEDGE_DETAIL, str(exc) or "Knowledge loading failed.")

    return _tool_success(ToolResultTypeSchema.KNOWLEDGE_DETAIL, _format_numbered_lines(first_line, lines))


@function_tool
async def find_knowledge(
    ctx: RunContextWrapper[AgentRuntimeContext],
    keyword: str,
    use_regex: bool = False,
    case_sensitive: bool = False,
    context_lines: int = _DEFAULT_FIND_CONTEXT_LINES,
    max_matches: int = _DEFAULT_FIND_MAX_MATCHES,
) -> str:
    """Search all knowledge bodies owned by the current agent for a keyword or regex.

    Args:
        keyword: str literal keyword or regular expression to search for across all current-agent knowledge bodies.
        use_regex: bool whether keyword should be interpreted as a Python regular expression.
        case_sensitive: bool whether matching should be case-sensitive.
        context_lines: int number of surrounding body lines to include before and after each match.
        max_matches: int maximum number of matching lines to return across all knowledge files.

    Returns:
        JSON status with grep-like knowledge name, body line numbers, and matching context.
    """
    try:
        agent_code = _valid_agent_code(ctx)
        needle = _valid_find_keyword(keyword)
        context = _valid_context_lines(context_lines)
        limit = _valid_max_matches(max_matches)
        matcher = _build_knowledge_matcher(needle, use_regex, case_sensitive)
        matches = _find_knowledge_matches(agent_code, matcher, context, limit)
    except Exception as exc:
        return _tool_error(ToolResultTypeSchema.KNOWLEDGE_DETAIL, str(exc) or "Knowledge search failed.")

    return _tool_success(ToolResultTypeSchema.KNOWLEDGE_DETAIL, _format_knowledge_matches(matches))


@function_tool
async def update_knowledge(
    ctx: RunContextWrapper[AgentRuntimeContext],
    target_name: str,
    name: str | None = None,
    description: str | None = None,
    start_line: int | None = None,
    end_line: int | None = None,
    content: str | None = None,
) -> str:
    """Update knowledge metadata and/or body lines for the current agent.

    Args:
        target_name: str current knowledge file name without .md.
        name: str | None optional new knowledge name. Updating it also renames the file.
        description: str | None optional new knowledge description.
        start_line: int | None optional 1-based body line to start replacing.
        end_line: int | None optional 1-based body line to stop replacing, inclusive.
        content: str | None replacement markdown body content for the specified body line range.

    Returns:
        JSON status describing the update result.
    """
    try:
        agent_code = _valid_agent_code(ctx)
        old_name = _valid_knowledge_name(target_name)
        new_name = _valid_knowledge_name(name) if name is not None else old_name
        new_description = _valid_description(description) if description is not None else None
        patch_body = _has_body_patch(start_line, end_line, content)
        if not patch_body and new_name == old_name and new_description is None:
            raise ValueError("At least one of name, description, or body line patch is required.")

        old_path = _existing_knowledge_path(agent_code, old_name)
        new_path = _safe_knowledge_path(agent_code, new_name)
        async with _locked_knowledge_paths(old_path, new_path):
            if new_path != old_path and new_path.exists():
                raise ValueError(f"Knowledge already exists for agent {agent_code}: {new_name}")
            doc = KnowledgeDocument.parse(old_path.read_text(encoding="utf-8"))
            doc.name = new_name
            if new_description is not None:
                doc.description = new_description
            if patch_body:
                doc.patch_body(start_line, end_line, _valid_content(content))
            new_path.parent.mkdir(parents=True, exist_ok=True)
            _write_text_atomic(new_path, doc.render())
            if new_path != old_path:
                old_path.unlink()
        _bump_knowledge_generation()
    except Exception as exc:
        return _tool_error(ToolResultTypeSchema.KNOWLEDGE_MUTATION, str(exc) or "Knowledge update failed.")

    changes = []
    if new_name != old_name:
        changes.append(f"renamed {old_name} -> {new_name}")
    if new_description is not None:
        changes.append("description updated")
    if patch_body:
        changes.append(f"body lines {start_line}-{end_line} replaced")
    return _tool_success(
        ToolResultTypeSchema.KNOWLEDGE_MUTATION,
        f"Knowledge updated for agent {agent_code}: {new_name}; " + "; ".join(changes),
    )


def load_knowledge_metadata(agent_code: str) -> tuple[str, ...]:
    if not _KNOWLEDGE_NAME_PATTERN.fullmatch(agent_code):
        return ()

    knowledges_dir = _agent_knowledges_dir(agent_code)
    if not knowledges_dir.is_dir():
        return ()

    blocks: list[str] = []
    knowledge_files = sorted(knowledges_dir.glob(f"*{KNOWLEDGE_EXTENSION}"), key=lambda path: path.stem)
    for knowledge_file in knowledge_files:
        knowledge_name = knowledge_file.stem
        if not _KNOWLEDGE_NAME_PATTERN.fullmatch(knowledge_name):
            continue
        if not _is_safe_knowledge_path(agent_code, knowledge_file):
            continue
        try:
            text = knowledge_file.read_text(encoding="utf-8")
            front_matter, body = _split_front_matter(text)
            metadata = yaml.safe_load(front_matter) if front_matter else None
        except (OSError, ValueError, yaml.YAMLError):
            continue
        if not isinstance(metadata, dict):
            continue
        metadata["body_line_count"] = len(body.split("\n")) if body else 0
        metadata = yaml.safe_dump(
            metadata,
            allow_unicode=True,
            default_flow_style=False,
            sort_keys=False,
        ).strip()
        blocks.append(f"### {knowledge_name}\n\n```yaml\n{metadata}\n```")
    return tuple(blocks)


def _valid_agent_code(ctx: RunContextWrapper[AgentRuntimeContext]) -> str:
    agent_code = ctx.context.agent_code.strip()
    if not _KNOWLEDGE_NAME_PATTERN.fullmatch(agent_code):
        raise ValueError("No valid agent code is available in runtime context.")
    return agent_code


def _valid_knowledge_name(name: str | None) -> str:
    knowledge_name = (name or "").strip()
    if not _KNOWLEDGE_NAME_PATTERN.fullmatch(knowledge_name):
        raise ValueError("Knowledge name must contain only letters, numbers, dot, underscore, or dash.")
    return knowledge_name


def _valid_description(description: str | None) -> str:
    value = (description or "").strip()
    if not value:
        raise ValueError("Knowledge description is required.")
    if len(value) > _MAX_DESCRIPTION_LENGTH:
        raise ValueError(f"Knowledge description must be at most {_MAX_DESCRIPTION_LENGTH} characters.")
    return value


def _valid_content(content: str | None) -> str:
    value = (content or "").strip("\r\n")
    if len(value) > _MAX_CONTENT_LENGTH:
        raise ValueError(f"Knowledge content must be at most {_MAX_CONTENT_LENGTH} characters.")
    return value


def _valid_find_keyword(keyword: str | None) -> str:
    value = keyword or ""
    if not value.strip():
        raise ValueError("Knowledge search keyword is required.")
    if "\n" in value or "\r" in value:
        raise ValueError("Knowledge search keyword must be a single line.")
    if len(value) > _MAX_FIND_KEYWORD_LENGTH:
        raise ValueError(f"Knowledge search keyword must be at most {_MAX_FIND_KEYWORD_LENGTH} characters.")
    return value


def _valid_context_lines(context_lines: int) -> int:
    if not isinstance(context_lines, int):
        raise ValueError("context_lines must be an integer.")
    if context_lines < 0 or context_lines > _MAX_FIND_CONTEXT_LINES:
        raise ValueError(f"context_lines must be between 0 and {_MAX_FIND_CONTEXT_LINES}.")
    return context_lines


def _valid_max_matches(max_matches: int) -> int:
    if not isinstance(max_matches, int):
        raise ValueError("max_matches must be an integer.")
    if max_matches < 1 or max_matches > _MAX_FIND_MAX_MATCHES:
        raise ValueError(f"max_matches must be between 1 and {_MAX_FIND_MAX_MATCHES}.")
    return max_matches


def _existing_knowledge_path(agent_code: str, knowledge_name: str) -> Path:
    knowledge_path = _safe_knowledge_path(agent_code, knowledge_name)
    if not knowledge_path.is_file():
        raise ValueError(f"Knowledge not found for agent {agent_code}: {knowledge_name}")
    return knowledge_path


def _safe_knowledge_path(agent_code: str, knowledge_name: str) -> Path:
    knowledge_path = _knowledge_entrypoint_path(agent_code, knowledge_name)
    if not _is_safe_knowledge_path(agent_code, knowledge_path):
        raise ValueError("Knowledge path is outside the current agent knowledges directory.")
    return knowledge_path


def _has_body_patch(start_line: int | None, end_line: int | None, content: str | None) -> bool:
    values = (start_line, end_line, content)
    if all(value is None for value in values):
        return False
    if any(value is None for value in values):
        raise ValueError("Body updates require start_line, end_line, and content.")
    return True


def _normalize_line_range(
    start_line: int | None,
    end_line: int | None,
    total_lines: int,
    *,
    allow_default: bool,
) -> tuple[int | None, int | None]:
    if start_line is None and end_line is None:
        if allow_default:
            return None, None
        raise ValueError("Body line range is required.")
    start = 1 if start_line is None and allow_default else start_line
    end = total_lines if end_line is None and allow_default else end_line
    if not isinstance(start, int) or not isinstance(end, int):
        raise ValueError("Body line range must be integer values.")
    if start < 1 or end < start:
        raise ValueError("Body line range is invalid.")
    if total_lines == 0:
        raise ValueError("Knowledge body is empty.")
    if end > total_lines:
        raise ValueError(f"Body line range exceeds knowledge body length ({total_lines}).")
    return start, end


def _format_numbered_lines(first_line: int, lines: list[str]) -> str:
    if not lines:
        return ""
    return "\n".join(f"{line_number} | {line}" for line_number, line in enumerate(lines, start=first_line))


def _build_knowledge_matcher(keyword: str, use_regex: bool, case_sensitive: bool) -> Callable[[str], bool]:
    if use_regex:
        flags = 0 if case_sensitive else regex.IGNORECASE
        try:
            pattern = regex.compile(keyword, flags)
        except regex.error as exc:
            raise ValueError(f"Knowledge search regex is invalid: {exc}") from exc

        def regex_match(line: str) -> bool:
            try:
                return pattern.search(line, timeout=_REGEX_SEARCH_TIMEOUT_SECONDS) is not None
            except TimeoutError as exc:
                raise ValueError("Knowledge search regex timed out.") from exc

        return regex_match

    needle = keyword if case_sensitive else keyword.casefold()

    def literal_match(line: str) -> bool:
        haystack = line if case_sensitive else line.casefold()
        return needle in haystack

    return literal_match


def _find_knowledge_matches(
    agent_code: str,
    matcher: Callable[[str], bool],
    context_lines: int,
    max_matches: int,
) -> list[tuple[str, int, list[tuple[int, str, bool]]]]:
    knowledges_dir = _agent_knowledges_dir(agent_code)
    if not knowledges_dir.is_dir():
        return []

    matches: list[tuple[str, int, list[tuple[int, str, bool]]]] = []
    matched_lines = 0
    knowledge_files = sorted(knowledges_dir.glob(f"*{KNOWLEDGE_EXTENSION}"), key=lambda path: path.stem)
    for knowledge_file in knowledge_files:
        knowledge_name = knowledge_file.stem
        if not _KNOWLEDGE_NAME_PATTERN.fullmatch(knowledge_name):
            continue
        if not _is_safe_knowledge_path(agent_code, knowledge_file):
            continue
        try:
            doc = KnowledgeDocument.parse(knowledge_file.read_text(encoding="utf-8"))
        except (OSError, ValueError):
            continue

        lines = doc.body_lines()
        for line_index, line in enumerate(lines, start=1):
            if not matcher(line):
                continue
            start = max(1, line_index - context_lines)
            end = min(len(lines), line_index + context_lines)
            context = [
                (number, lines[number - 1], number == line_index)
                for number in range(start, end + 1)
            ]
            matches.append((knowledge_name, line_index, context))
            matched_lines += 1
            if matched_lines >= max_matches:
                return matches
    return matches


def _format_knowledge_matches(matches: list[tuple[str, int, list[tuple[int, str, bool]]]]) -> str:
    if not matches:
        return "No matching knowledge lines found."

    blocks: list[str] = []
    for knowledge_name, matched_line, context in matches:
        lines = [f"## {knowledge_name}:{matched_line}"]
        for line_number, line, is_match in context:
            marker = ">" if is_match else " "
            lines.append(f"{marker} {line_number} | {line}")
        blocks.append("\n".join(lines))
    return "\n\n".join(blocks)


def _split_front_matter(text: str) -> tuple[str, str]:
    lines = text.splitlines(keepends=True)
    if not lines or lines[0].strip() != "---":
        raise ValueError("Knowledge file must start with YAML Front Matter.")
    for index, line in enumerate(lines[1:], start=1):
        if line.strip() == "---":
            front_matter = "".join(lines[1:index])
            body = "".join(lines[index + 1:]).lstrip("\r\n").rstrip("\n")
            return front_matter, body
    raise ValueError("Knowledge file has unterminated YAML Front Matter.")


def _write_text_atomic(path: Path, text: str) -> None:
    temp_name = ""
    try:
        with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as tmp:
            temp_name = tmp.name
            tmp.write(text)
        Path(temp_name).replace(path)
    finally:
        if temp_name:
            try:
                Path(temp_name).unlink(missing_ok=True)
            except OSError:
                pass


@asynccontextmanager
async def _locked_knowledge_paths(*paths: Path):
    locks = [_knowledge_locks.setdefault(str(path.resolve()), asyncio.Lock()) for path in sorted(set(paths), key=str)]
    for lock in locks:
        await lock.acquire()
    try:
        yield
    finally:
        for lock in reversed(locks):
            lock.release()


def _tool_success(result_type: ToolResultTypeSchema, output: str) -> str:
    return ToolResultSchema(
        status=ToolResultStatusSchema.SUCCESS,
        type=result_type,
        output=output,
    ).model_dump_json()


def _tool_error(result_type: ToolResultTypeSchema, output: str) -> str:
    return ToolResultSchema(
        status=ToolResultStatusSchema.ERROR,
        type=result_type,
        output=output,
    ).model_dump_json()


def _knowledge_entrypoint_path(agent_code: str, knowledge_name: str) -> Path:
    return _agent_knowledges_dir(agent_code) / f"{knowledge_name}{KNOWLEDGE_EXTENSION}"


def _agent_knowledges_dir(agent_code: str) -> Path:
    return WORKSPACE / "agents" / agent_code / KNOWLEDGES_DIR_NAME


def _is_safe_knowledge_path(agent_code: str, path: Path) -> bool:
    try:
        return path.resolve().is_relative_to(_agent_knowledges_dir(agent_code).resolve())
    except OSError:
        return False
