"""Test Generation Agent — auto-generates pytest stubs for untested code."""
from __future__ import annotations

import ast
import logging
from pathlib import Path
from typing import Any, Dict, List, Set

from mesh.codex.base_agent import CodexAgent
from mesh.codex.models import CodexTask, CodexTaskResult, FindingSeverity

logger = logging.getLogger(__name__)

# Project root for relative path computation
_PROJECT_ROOT = Path(__file__).resolve().parents[2]


class TestGeneratorAgent(CodexAgent):
    """Scans source files, identifies untested functions/classes, and
    generates pytest test stubs with reasonable assertions.

    Writes generated tests to tests/generated/ and respects existing
    test patterns from tests/conftest.py.
    """

    def __init__(self, mesh_db, config=None):
        super().__init__("test_gen", mesh_db, config)

    async def execute(self, task: CodexTask) -> CodexTaskResult:
        target = task.target or str(_PROJECT_ROOT / "mesh")
        source_files = self._collect_python_files(target)
        if not source_files:
            return CodexTaskResult(
                task_id=task.task_id,
                success=True,
                findings=[],
                recommendations=["No Python source files found at target"],
            )

        # Find existing tests
        test_dir = _PROJECT_ROOT / "tests"
        existing_tests = self._collect_python_files(str(test_dir))
        tested_functions = self._extract_tested_functions(existing_tests)

        # Scan source for untested functions
        findings: List[Dict[str, Any]] = []
        generated_files: Dict[str, str] = {}
        total_functions = 0
        untested_count = 0

        for filepath in source_files:
            content = self._get_context(str(filepath))
            if content is None:
                continue

            try:
                tree = ast.parse(content, filename=str(filepath))
            except SyntaxError:
                continue

            functions = self._extract_public_symbols(tree)
            total_functions += len(functions)

            untested = [f for f in functions if f["name"] not in tested_functions]
            untested_count += len(untested)

            if untested:
                findings.append({
                    "severity": FindingSeverity.INFO.value,
                    "category": "coverage_gap",
                    "location": str(filepath),
                    "message": f"{len(untested)}/{len(functions)} public functions untested",
                    "suggestion": ", ".join(f["name"] for f in untested[:10]),
                })

                # Generate test file
                test_content = self._generate_test_file(filepath, untested, content)
                if test_content:
                    # Compute output path
                    rel = filepath.relative_to(_PROJECT_ROOT)
                    test_filename = f"test_{rel.stem}.py"
                    test_path = _PROJECT_ROOT / "tests" / "generated" / test_filename
                    generated_files[str(test_path)] = test_content

        # Write generated tests if requested
        artifacts: Dict[str, str] = {}
        write_tests = task.parameters.get("write", True)
        if write_tests and generated_files:
            gen_dir = _PROJECT_ROOT / "tests" / "generated"
            gen_dir.mkdir(parents=True, exist_ok=True)

            # Write __init__.py if missing
            init_file = gen_dir / "__init__.py"
            if not init_file.exists():
                init_file.write_text('"""Auto-generated tests from Codex test_gen agent."""\n')

            for path, content in generated_files.items():
                Path(path).write_text(content)
                artifacts[Path(path).name] = path
                logger.info("Generated test file: %s", path)

        recommendations = []
        if total_functions > 0:
            coverage_pct = ((total_functions - untested_count) / total_functions) * 100
            recommendations.append(
                f"Coverage estimate: {coverage_pct:.0f}% ({total_functions - untested_count}/{total_functions} functions have tests)"
            )
        if untested_count > 0:
            recommendations.append(
                f"Generated test stubs for {untested_count} untested functions in {len(generated_files)} files"
            )
        if artifacts:
            recommendations.append(
                f"Test files written to tests/generated/ — review and add assertions as needed"
            )

        return CodexTaskResult(
            task_id=task.task_id,
            success=True,
            findings=findings,
            recommendations=recommendations,
            artifacts=artifacts,
        )

    def _extract_tested_functions(self, test_files: List[Path]) -> Set[str]:
        """Extract function names that appear to be tested based on test function names."""
        tested = set()
        for filepath in test_files:
            content = self._get_context(str(filepath))
            if content is None:
                continue
            try:
                tree = ast.parse(content)
            except SyntaxError:
                continue

            for node in ast.walk(tree):
                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    name = node.name
                    # test_foo_bar -> foo_bar
                    if name.startswith("test_"):
                        tested.add(name[5:])
                    # Also track direct references in the test body
                    for child in ast.walk(node):
                        if isinstance(child, ast.Name):
                            tested.add(child.id)
                        elif isinstance(child, ast.Attribute):
                            tested.add(child.attr)
        return tested

    def _extract_public_symbols(self, tree: ast.AST) -> List[Dict[str, Any]]:
        """Extract public functions and classes from an AST."""
        symbols = []
        for node in ast.iter_child_nodes(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                if not node.name.startswith("_"):
                    symbols.append({
                        "name": node.name,
                        "type": "async_function" if isinstance(node, ast.AsyncFunctionDef) else "function",
                        "lineno": node.lineno,
                        "args": [arg.arg for arg in node.args.args if arg.arg != "self"],
                        "has_return": self._has_return(node),
                        "decorators": [self._decorator_name(d) for d in node.decorator_list],
                    })
            elif isinstance(node, ast.ClassDef):
                if not node.name.startswith("_"):
                    methods = []
                    for child in node.body:
                        if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
                            if not child.name.startswith("_") or child.name == "__init__":
                                methods.append({
                                    "name": child.name,
                                    "args": [a.arg for a in child.args.args if a.arg != "self"],
                                })
                    symbols.append({
                        "name": node.name,
                        "type": "class",
                        "lineno": node.lineno,
                        "methods": methods,
                    })
        return symbols

    def _has_return(self, node: ast.AST) -> bool:
        for child in ast.walk(node):
            if isinstance(child, ast.Return) and child.value is not None:
                return True
        return False

    def _decorator_name(self, node: ast.AST) -> str:
        if isinstance(node, ast.Name):
            return node.id
        if isinstance(node, ast.Attribute):
            return node.attr
        if isinstance(node, ast.Call):
            return self._decorator_name(node.func)
        return ""

    def _generate_test_file(self, source_path: Path, untested: List[Dict[str, Any]],
                            source_content: str) -> str:
        """Generate a pytest test file for untested symbols."""
        rel_path = source_path.relative_to(_PROJECT_ROOT)
        module_path = ".".join(rel_path.with_suffix("").parts)

        lines = [
            f'"""Auto-generated tests for {module_path}.',
            '',
            'Generated by Codex test_gen agent. Review and add meaningful assertions.',
            '"""',
            'from __future__ import annotations',
            '',
            'import pytest',
            '',
        ]

        # Build import line
        names_to_import = []
        for sym in untested:
            if sym["type"] in ("function", "async_function", "class"):
                names_to_import.append(sym["name"])

        if names_to_import:
            # Use try/except to handle import failures gracefully
            lines.append(f'try:')
            lines.append(f'    from {module_path} import (')
            for name in names_to_import:
                lines.append(f'        {name},')
            lines.append(f'    )')
            lines.append(f'except ImportError:')
            lines.append(f'    pytest.skip("Cannot import {module_path}", allow_module_level=True)')
            lines.append('')
            lines.append('')

        # Generate test functions
        for sym in untested:
            if sym["type"] == "function":
                lines.extend(self._gen_function_test(sym))
            elif sym["type"] == "async_function":
                lines.extend(self._gen_async_function_test(sym))
            elif sym["type"] == "class":
                lines.extend(self._gen_class_test(sym))

        return "\n".join(lines)

    def _gen_function_test(self, sym: Dict[str, Any]) -> List[str]:
        name = sym["name"]
        args = sym.get("args", [])
        has_return = sym.get("has_return", False)

        lines = [
            f'def test_{name}():',
            f'    """Test {name}."""',
        ]

        if args:
            # Generate placeholder arguments
            arg_str = ", ".join(f'{a}=None' for a in args)
            if has_return:
                lines.append(f'    result = {name}({arg_str})')
                lines.append(f'    assert result is not None  # TODO: add meaningful assertion')
            else:
                lines.append(f'    {name}({arg_str})  # TODO: verify side effects')
        else:
            if has_return:
                lines.append(f'    result = {name}()')
                lines.append(f'    assert result is not None  # TODO: add meaningful assertion')
            else:
                lines.append(f'    {name}()  # TODO: verify side effects')

        lines.append('')
        lines.append('')
        return lines

    def _gen_async_function_test(self, sym: Dict[str, Any]) -> List[str]:
        name = sym["name"]
        args = sym.get("args", [])
        has_return = sym.get("has_return", False)

        lines = [
            '@pytest.mark.asyncio',
            f'async def test_{name}():',
            f'    """Test {name}."""',
        ]

        if args:
            arg_str = ", ".join(f'{a}=None' for a in args)
            if has_return:
                lines.append(f'    result = await {name}({arg_str})')
                lines.append(f'    assert result is not None  # TODO: add meaningful assertion')
            else:
                lines.append(f'    await {name}({arg_str})  # TODO: verify side effects')
        else:
            if has_return:
                lines.append(f'    result = await {name}()')
                lines.append(f'    assert result is not None  # TODO: add meaningful assertion')
            else:
                lines.append(f'    await {name}()  # TODO: verify side effects')

        lines.append('')
        lines.append('')
        return lines

    def _gen_class_test(self, sym: Dict[str, Any]) -> List[str]:
        name = sym["name"]
        methods = sym.get("methods", [])

        lines = [
            f'class Test{name}:',
            f'    """Tests for {name}."""',
            '',
        ]

        # Test instantiation
        init_method = next((m for m in methods if m["name"] == "__init__"), None)
        if init_method:
            init_args = [a for a in init_method.get("args", []) if a != "self"]
            if init_args:
                arg_str = ", ".join(f'{a}=None' for a in init_args)
                lines.append(f'    def test_init(self):')
                lines.append(f'        obj = {name}({arg_str})')
                lines.append(f'        assert obj is not None')
            else:
                lines.append(f'    def test_init(self):')
                lines.append(f'        obj = {name}()')
                lines.append(f'        assert obj is not None')
        else:
            lines.append(f'    def test_init(self):')
            lines.append(f'        obj = {name}()')
            lines.append(f'        assert obj is not None')

        lines.append('')

        # Test each public method
        for method in methods:
            if method["name"] == "__init__":
                continue
            mname = method["name"]
            margs = [a for a in method.get("args", []) if a != "self"]
            lines.append(f'    def test_{mname}(self):')
            lines.append(f'        """Test {name}.{mname}."""')
            if init_method:
                init_args = [a for a in init_method.get("args", []) if a != "self"]
                arg_str = ", ".join(f'{a}=None' for a in init_args) if init_args else ""
                lines.append(f'        obj = {name}({arg_str})')
            else:
                lines.append(f'        obj = {name}()')
            if margs:
                call_args = ", ".join(f'{a}=None' for a in margs)
                lines.append(f'        obj.{mname}({call_args})  # TODO: add assertions')
            else:
                lines.append(f'        obj.{mname}()  # TODO: add assertions')
            lines.append('')

        lines.append('')
        return lines
