from __future__ import annotations

import asyncio
import dataclasses
import glob
import json
import logging
import os
import subprocess
import time
import uuid
from collections.abc import Callable
from pathlib import Path
from typing import Any

from py_agent.config import PyAgentConfig
from py_agent.infrastructure.auth_service import get_jwt_token
from py_agent.infrastructure.backend_client import (
    post_update_operation_metrics,
)
from py_agent.infrastructure.coverage_runner import convert_coverage_to_json, run_coverage
from py_agent.infrastructure.pytest_runner import run_pytest
from py_agent.libs.queue import Queue
from py_agent.models.testable import Testable
from py_agent.services.agent_runner import run_agent
from py_agent.services.command_discovery import (
    discover_format_command,
    discover_lint_command,
    discover_test_command,
)
from py_agent.services.prompt_builder import (
    SYSTEM_PROMPT,
    _adjust_pythonpath_for_relative_imports,
    build_user_prompt,
)
from py_agent.services.testables_extractor import TestablesExtractor
from py_agent.services.tested_code_extractor import TestedCodeExtractor
from py_agent.utils.pytest_config import read_pytest_config
from py_agent.utils.test_file_utils import compute_test_file_path, create_conftest, create_scaffold

logger = logging.getLogger(__name__)


# ── Types (mirrors ts-agent's generate-tests.types.ts) ──────────────────────


@dataclasses.dataclass
class GenerateTestsQueueData:
    file_path: str
    method_type: str
    method_name: str
    parent_name: str = ""


@dataclasses.dataclass
class TestableCompleteResult:
    test_file_path: str
    request_id: str
    cost_usd: float | None = None
    tests_count: int = 0


# Mirrors ts-agent's TestableCompleteCallback:
#   (testable_item, result) where testable_item = {"file_path": str, "testable": Testable}
TestableCompleteCallback = Callable[
    [dict[str, Any], TestableCompleteResult | None],
    None,
]

TestableStartCallback = Callable[[dict[str, Any]], None]


# ── PyAgent ──────────────────────────────────────────────────────────────────


class PyAgent:
    def __init__(self, config: PyAgentConfig | None = None):
        self._config = config or PyAgentConfig()
        self._initialized = False
        self._queue: Queue[GenerateTestsQueueData] | None = None
        self._test_results: dict[str, TestableCompleteResult | None] = {}

    # ── Queue ────────────────────────────────────────────────────────────────

    @staticmethod
    def _get_queue_item_key(data: GenerateTestsQueueData) -> str:
        return f"{data.file_path}:{data.method_type}:{data.parent_name}:{data.method_name}"

    def _get_queue(self) -> Queue[GenerateTestsQueueData]:
        if self._queue is None:
            self._queue = Queue(
                concurrency=self._config.concurrency,
                get_item_key=self._get_queue_item_key,
            )
        return self._queue

    # ── Lifecycle ────────────────────────────────────────────────────────────

    async def init(self, api_key: str | None = None) -> None:
        if api_key:
            self._config.api_key = api_key
        self._initialized = True

    # ── Testables ────────────────────────────────────────────────────────────

    async def get_testables(
        self, pattern: str, project_path: str | Path | None = None
    ) -> list[Testable]:
        project_root = str(project_path) if project_path else os.getcwd()
        matched_files = glob.glob(pattern, recursive=True)

        all_testables: list[Testable] = []
        for file_path in matched_files:
            abs_path = os.path.abspath(file_path)
            try:
                rel_path = os.path.relpath(abs_path, project_root)
            except ValueError:
                rel_path = abs_path

            try:
                testables_dto = TestablesExtractor.get_testables(abs_path, rel_path, project_root)
                all_testables.extend(testables_dto.testables)
            except Exception:
                logger.exception(f"Failed to extract testables from {abs_path}")

        return all_testables

    # ── Core generation logic ─────────────────────────────────────────────────

    @staticmethod
    def _get_git_url(project_root: str) -> str | None:
        """Detect git remote URL from the project root."""
        try:
            result = subprocess.run(
                ["git", "-C", project_root, "remote", "get-url", "origin"],
                capture_output=True,
                text=True,
                timeout=5,
            )
            if result.returncode == 0 and result.stdout.strip():
                return result.stdout.strip()
        except Exception:
            pass
        return None

    @staticmethod
    def _build_final_test_result(
        pytest_result: dict[str, Any],
        test_file_path: str,
        rel_test_path: str,
    ) -> dict[str, Any] | None:
        """Build finalTestResult matching backend testResultDataSourceSchema."""
        report = pytest_result.get("pytestReport")
        command_str = pytest_result.get("command", "")

        # Read the full generated test file
        test_code = ""
        try:
            with open(test_file_path, encoding="utf-8") as f:
                test_code = f.read()
        except Exception:
            pass

        if report is None:
            # pytest crashed or XML couldn't be parsed — still report what we can
            stderr = pytest_result.get("stderr", "")
            return {
                "describe": {
                    "code": test_code,
                    "path": rel_test_path,
                    "status": "RuntimeError",
                    "tests": [],
                    "lintErrors": [],
                },
                "stdout": pytest_result.get("stdout", ""),
                "stderr": stderr,
                "exitCode": -1,
                "command": command_str,
                "testsCount": 0,
                "passedCount": 0,
                "failedCount": 0,
                "failureReason": stderr[:2000] if stderr else "pytest report could not be parsed",
            }

        # Map individual test results
        tests: list[dict[str, Any]] = []
        for t in report.get("tests", []):
            status = "Green" if t.get("outcome") == "passed" else "Red"
            test_entry: dict[str, Any] = {
                "code": "",
                "name": t.get("name", ""),
                "status": status,
            }
            if t.get("outcome") == "failed":
                errors = []
                if t.get("message"):
                    errors.append(
                        {
                            "shortDescription": t["message"],
                            "fullDescription": t.get("longrepr", t["message"]),
                        }
                    )
                test_entry["errors"] = errors
            tests.append(test_entry)

        # Determine overall describe status
        summary = report.get("summary", {})
        has_failures = summary.get("failed", 0) > 0 or summary.get("error", 0) > 0
        describe_status = "Grey" if has_failures else "Green"

        exit_code = report.get("exitCode")
        stderr = pytest_result.get("stderr", "")

        result: dict[str, Any] = {
            "describe": {
                "code": test_code,
                "path": rel_test_path,
                "status": describe_status,
                "tests": tests,
                "lintErrors": [],
            },
            "stdout": pytest_result.get("stdout", ""),
            "stderr": stderr,
            "exitCode": exit_code,
            "command": command_str,
            "testsCount": summary.get("total", 0),
            "passedCount": summary.get("passed", 0),
            "failedCount": summary.get("failed", 0),
        }

        # Only include failureReason when pytest itself failed (not just test failures)
        if exit_code is not None and exit_code > 1:
            result["failureReason"] = (
                stderr[:2000] if stderr else f"pytest exited with code {exit_code}"
            )

        return result

    @staticmethod
    def _build_test_counters(
        pytest_result: dict[str, Any],
    ) -> tuple[int, dict[str, int]]:
        """Extract finalGreenTestsCount and finalTestCounters from pytest result."""
        report = pytest_result.get("pytestReport")
        if report is None:
            return 0, {
                "greenTestsCount": 0,
                "greyTestsCount": 0,
                "redTestsCount": 0,
                "naTestsCount": 0,
                "errorSuitsCount": 0,
                "lintErrorsCount": 0,
                "whiteTestsCount": 0,
            }
        summary = report.get("summary", {})
        green = summary.get("passed", 0)
        return green, {
            "greenTestsCount": green,
            "greyTestsCount": 0,
            "redTestsCount": summary.get("failed", 0),
            "naTestsCount": 0,
            "errorSuitsCount": summary.get("error", 0),
            "lintErrorsCount": 0,
            "whiteTestsCount": 0,
        }

    async def _generate_tests_impl(
        self,
        file_path: str | Path,
        testable: Testable,
        project_path: str | Path | None = None,
        user_prompt: str = "",
    ) -> dict[str, Any]:
        t_start = time.monotonic()

        abs_file_path = str(os.path.abspath(file_path))
        project_root = str(project_path) if project_path else os.getcwd()
        try:
            rel_path = os.path.relpath(abs_file_path, project_root)
        except ValueError:
            rel_path = abs_file_path

        class_name = testable.parent_name if testable.type == "method" else None

        # 1. Extract metadata (measure metadataCollectionTime)
        t_meta_start = time.monotonic()
        tested_code_dto = TestedCodeExtractor.get_tested_code_by_function_name(
            absolute_file_path=abs_file_path,
            relative_file_path=rel_path,
            project_path=project_root,
            function_name=testable.name,
            class_name=class_name,
            extraction_depth=self._config.dependency_extraction_depth,
        )
        metadata_collection_time = int((time.monotonic() - t_meta_start) * 1000)

        # 2. Compute test file path and create scaffold
        test_file_path = compute_test_file_path(
            source_file_path=abs_file_path,
            project_path=project_root,
            method_name=testable.name,
            test_structure=self._config.test_structure,
            test_file_name=self._config.test_file_name,
            class_name=class_name,
        )
        create_scaffold(test_file_path, project_path=project_root)

        rel_test_path = os.path.relpath(test_file_path, project_root)

        # Detect git URL
        git_url = self._get_git_url(project_root)

        # 3. Authenticate and get JWT
        if not self._config.api_key:
            raise RuntimeError(
                "api_key is required for test generation. Call init(api_key=...) first."
            )
        jwt_token = await get_jwt_token(self._config.api_base_url, self._config.api_key)

        # 4. Build prompt locally (bypasses backend /generate-prompt)
        request_id = str(uuid.uuid4())
        pytest_config = read_pytest_config(project_root)
        file_code = (
            tested_code_dto.extracted_data.file_code if tested_code_dto.extracted_data else None
        )
        effective_pythonpath, _ = _adjust_pythonpath_for_relative_imports(
            pytest_config.pythonpath, project_root, file_code
        )
        create_conftest(test_file_path, project_root, effective_pythonpath)
        pythonpath_prefix = ""
        if effective_pythonpath != list(pytest_config.pythonpath):
            pythonpath_prefix = f"PYTHONPATH={':'.join(effective_pythonpath)}:$PYTHONPATH "
        test_command = discover_test_command(project_root, rel_test_path, pythonpath_prefix)
        lint_command = discover_lint_command(project_root, rel_test_path)
        format_command = discover_format_command(project_root, rel_test_path)
        built_user_prompt = build_user_prompt(
            tested_code=tested_code_dto,
            test_file_path=test_file_path,
            rel_test_path=rel_test_path,
            rel_source_path=rel_path,
            project_root=project_root,
            pytest_config=pytest_config,
            test_command=test_command,
            lint_command=lint_command,
            format_command=format_command,
        )
        system_prompt = SYSTEM_PROMPT

        # 4b. Create initial operation metrics row for monitoring
        full_prompt = f"## System Prompt\n{system_prompt}\n\n## User Prompt\n{built_user_prompt}"
        if os.environ.get("EARLY_SHOW_PROMPT", "").lower() in ("1", "true", "yes"):
            print(full_prompt)

        await post_update_operation_metrics(
            api_base_url=self._config.api_base_url,
            jwt_token=jwt_token,
            request_id=request_id,
            request_source=self._config.request_source,
            git_url=git_url,
            prompt=full_prompt,
            workflow_run_id=self._config.workflow_run_id,
            file_path=rel_path,
            testable_name=testable.name,
            class_name=class_name,
            language="python",
            test_framework="pytest",
            metadata_collection_time=metadata_collection_time,
        )

        # 5. Run Claude Agent SDK (measure responseProcessingTime)
        t_sdk_start = time.monotonic()
        result = await run_agent(
            user_prompt=built_user_prompt,
            system_prompt=system_prompt,
            test_file_path=test_file_path,
            working_directory=project_root,
            jwt_token=jwt_token,
            api_base_url=self._config.api_base_url,
            request_id=request_id,
            testable_name=testable.name,
        )
        response_processing_time = int((time.monotonic() - t_sdk_start) * 1000)

        # 6. Run lint --fix + formatter on the generated test file (post-agent safety net)
        if lint_command:
            try:
                subprocess.run(
                    lint_command,
                    shell=True,
                    cwd=project_root,
                    capture_output=True,
                    timeout=30,
                )
            except Exception:
                logger.debug("Post-generation lint command failed (non-fatal)")

        if format_command:
            try:
                subprocess.run(
                    format_command,
                    shell=True,
                    cwd=project_root,
                    capture_output=True,
                    timeout=30,
                )
            except Exception:
                logger.debug("Post-generation format command failed (non-fatal)")

        # 6b. Prepend "# mypy: ignore-errors" so strict client mypy configs
        #     don't fail on generated test files.
        try:
            with open(test_file_path, encoding="utf-8") as f:
                content = f.read()
            if not content.startswith("# mypy: ignore-errors"):
                with open(test_file_path, "w", encoding="utf-8") as f:
                    f.write("# mypy: ignore-errors\n" + content)
        except Exception:
            logger.debug("Failed to prepend mypy ignore comment (non-fatal)")

        # 7. Use test counts from the agent's own pytest runs (via prune hook)
        agent_passed = result.get("passed", 0)
        agent_failed = result.get("failed", 0)
        final_green_tests_count = agent_passed
        final_test_counters: dict[str, int] = {
            "greenTestsCount": agent_passed,
            "greyTestsCount": 0,
            "redTestsCount": agent_failed,
            "naTestsCount": 0,
            "errorSuitsCount": 0,
            "lintErrorsCount": 0,
            "whiteTestsCount": 0,
        }
        final_test_result: dict[str, Any] | None = None

        if final_green_tests_count == 0:
            # Mirror ts-agent: delete scaffold file and report error metrics
            try:
                os.unlink(test_file_path)
            except OSError:
                logger.debug("Could not delete test file for failed generation: %s", test_file_path)

            time_elapsed = int((time.monotonic() - t_start) * 1000)
            error_detail = f"0 passing tests for {testable.name}"
            await post_update_operation_metrics(
                api_base_url=self._config.api_base_url,
                jwt_token=jwt_token,
                request_id=request_id,
                request_source=self._config.request_source,
                time_elapsed=time_elapsed,
                metadata_collection_time=metadata_collection_time,
                response_processing_time=response_processing_time,
                git_url=git_url,
                workflow_run_id=self._config.workflow_run_id,
                file_path=rel_path,
                testable_name=testable.name,
                class_name=class_name,
                language="python",
                test_framework="pytest",
                error_cause=error_detail,
            )
            raise RuntimeError(f"Test generation produced {error_detail}")

        time_elapsed = int((time.monotonic() - t_start) * 1000)

        # 8. Report operation metrics (mirrors ts-agent handleSdkResult)
        # Token counts and cost are NOT sent here — the backend proxy handles
        # all token/cost tracking via updateTokensCountByRequestId using the
        # x-request-id header passed through ANTHROPIC_CUSTOM_HEADERS.
        full_prompt = f"## System Prompt\n{system_prompt}\n\n## User Prompt\n{built_user_prompt}"
        await post_update_operation_metrics(
            api_base_url=self._config.api_base_url,
            jwt_token=jwt_token,
            request_id=request_id,
            request_source=self._config.request_source,
            final_green_tests_count=final_green_tests_count,
            final_test_counters=final_test_counters,
            time_elapsed=time_elapsed,
            metadata_collection_time=metadata_collection_time,
            response_processing_time=response_processing_time,
            git_url=git_url,
            final_test_result=final_test_result,
            prompt=full_prompt,
            workflow_run_id=self._config.workflow_run_id,
            file_path=rel_path,
            testable_name=testable.name,
            class_name=class_name,
            language="python",
            test_framework="pytest",
        )

        return {
            "test_file_path": test_file_path,
            "request_id": request_id,
            "cost_usd": result.get("cost_usd"),
            "tests_count": final_green_tests_count,
        }

    # ── Queue dispatch (mirrors ts-agent's addGenerationToQueue) ─────────────

    def _enqueue_generation(
        self,
        file_path: str | Path,
        testable: Testable,
        project_path: str | Path | None = None,
        user_prompt: str = "",
        progress: dict[str, int] | None = None,
    ) -> asyncio.Task[None] | None:
        """
        Synchronously submits a generation job to the queue.
        Returns the asyncio.Task (await it to wait for completion), or None if already queued.
        """
        abs_file_path = str(os.path.abspath(file_path))
        process_data = GenerateTestsQueueData(
            file_path=abs_file_path,
            method_type=testable.type,
            method_name=testable.name,
            parent_name=testable.parent_name or "",
        )
        queue = self._get_queue()

        if queue.has(process_data):
            logger.info(f'Test generation for "{testable.name}" is already in progress')
            return None

        if progress:
            logger.info(
                f"Generating tests for {testable.name} — "
                f"{progress['index']} of {progress['total']} functions"
            )
        else:
            logger.info(f"Generating tests for {testable.name} in {abs_file_path}")

        key = self._get_queue_item_key(process_data)

        async def task() -> None:
            try:
                result = await self._generate_tests_impl(
                    file_path=file_path,
                    testable=testable,
                    project_path=project_path,
                    user_prompt=user_prompt,
                )
                self._test_results[key] = TestableCompleteResult(
                    test_file_path=result["test_file_path"],
                    request_id=result["request_id"],
                    cost_usd=result.get("cost_usd"),
                    tests_count=result.get("tests_count", 0),
                )
            except Exception:
                logger.exception(f"Test generation failed for {testable.name}")
                self._test_results[key] = None

        return queue.enqueue(process_data, task)

    # ── Public API ────────────────────────────────────────────────────────────

    async def generate_tests(
        self,
        file_path: str | Path,
        testable: Testable,
        project_path: str | Path | None = None,
        user_prompt: str = "",
    ) -> dict[str, Any]:
        """Generate tests for a single testable, dispatched through the queue."""
        abs_file_path = str(os.path.abspath(file_path))
        key = self._get_queue_item_key(
            GenerateTestsQueueData(
                file_path=abs_file_path,
                method_type=testable.type,
                method_name=testable.name,
                parent_name=testable.parent_name or "",
            )
        )

        t = self._enqueue_generation(
            file_path=file_path,
            testable=testable,
            project_path=project_path,
            user_prompt=user_prompt,
        )
        if t is not None:
            await t

        result = self._test_results.pop(key, None)
        if result is None:
            raise RuntimeError(f"Test generation failed for {testable.name}")
        return {
            "test_file_path": result.test_file_path,
            "request_id": result.request_id,
            "cost_usd": result.cost_usd,
        }

    async def bulk_generate_tests(
        self,
        testables_to_generate: list[dict[str, Any]],
        on_testable_complete: TestableCompleteCallback | None = None,
        on_testable_start: TestableStartCallback | None = None,
        user_prompt: str = "",
    ) -> None:
        """
        Generate tests for multiple testables concurrently via the queue.

        Each entry must be {"file_path": str, "testable": Testable}.
        on_testable_start fires when a testable begins executing (acquires queue slot).
        on_testable_complete fires after each testable finishes (mirrors ts-agent).
        """
        queue = self._get_queue()

        def _find_item(key: str) -> dict[str, Any] | None:
            return next(
                (
                    i
                    for i in testables_to_generate
                    if self._get_queue_item_key(
                        GenerateTestsQueueData(
                            file_path=str(os.path.abspath(i["file_path"])),
                            method_type=i["testable"].type,
                            method_name=i["testable"].name,
                            parent_name=i["testable"].parent_name or "",
                        )
                    )
                    == key
                ),
                None,
            )

        if on_testable_start:

            def _on_start(key: str) -> None:
                item = _find_item(key)
                if item is not None:
                    try:
                        on_testable_start(item)
                    except Exception:
                        logger.exception(f"Start callback failed for {key}")

            queue.on_start(_on_start)

        if on_testable_complete:

            def _on_done(key: str) -> None:
                item = _find_item(key)
                if item is not None:
                    result = self._test_results.pop(key, None)
                    try:
                        on_testable_complete(item, result)
                    except Exception:
                        logger.exception(f"Completion callback failed for {key}")

            queue.on_done(_on_done)

        for index, item in enumerate(testables_to_generate):
            self._enqueue_generation(
                file_path=item["file_path"],
                testable=item["testable"],
                project_path=item.get("project_path"),
                user_prompt=user_prompt,
                progress={"index": index + 1, "total": len(testables_to_generate)},
            )

        await queue.get_idle_promise()

    # ── Test running & coverage ───────────────────────────────────────────────

    async def run_tests(
        self, test_file_path: str | Path, project_path: str | Path | None = None
    ) -> dict[str, Any]:
        abs_test_path = str(os.path.abspath(test_file_path))
        project_root = str(project_path) if project_path else os.getcwd()
        return run_pytest(test_file_path=abs_test_path, repository_root_path=project_root)

    async def generate_coverage(
        self,
        project_path: str | Path | None = None,
        coverage_report_path: str | Path | None = None,
    ) -> dict[str, Any]:
        project_root = str(project_path) if project_path else os.getcwd()
        report_path = (
            str(coverage_report_path)
            if coverage_report_path
            else os.path.join(project_root, ".coverage_report", "coverage.json")
        )
        return run_coverage(project_path=project_root, coverage_report_path=report_path)

    async def get_coverage_tree(
        self,
        coverage_report_path: str | Path,
        project_path: str | Path | None = None,
    ) -> dict[str, Any]:
        report_path = str(coverage_report_path)
        if not os.path.exists(report_path):
            return {"error": f"Coverage report not found at {report_path}"}

        with open(report_path, encoding="utf-8") as f:
            result: dict[str, Any] = json.load(f)
            return result

    async def convert_coverage(
        self,
        coverage_db_path: str | Path,
        coverage_report_path: str | Path,
        project_path: str | Path | None = None,
    ) -> dict[str, Any]:
        return convert_coverage_to_json(
            coverage_db_path=str(coverage_db_path),
            coverage_report_path=str(coverage_report_path),
        )

    # ── Cache / context ───────────────────────────────────────────────────────

    def on_file_changed(self, file_path: str | Path | None = None) -> None:
        from py_agent.services.import_jedi_extractor import JediContainer

        JediContainer.clear_all_caches()

    def update_context(self, context: dict[str, object]) -> None:
        for key, value in context.items():
            if hasattr(self._config, key):
                setattr(self._config, key, value)
