"""
py-agent CLI entry point.

Invoked by early-cli as a subprocess for Python projects.
Streams JSON-lines progress events to stdout so the caller can track progress.

Event schema:
    {"type": "scan",     "files": <int>, "target": <str>}
    {"type": "extract",  "file": <str>, "functions": <int>, "total_in_file": <int>, "skipped": <int|absent>, "changed": <int|absent>, "unchanged": <int|absent>}
    {"type": "start",    "total": <int>, "files": <int>}
    {"type": "skip",     "method": <str>, "file": <str>}
    {"type": "progress", "method": <str>, "file": <str>, "index": <int>, "total": <int>}
    {"type": "done",     "method": <str>, "file": <str>, "test_file": <str>, "request_id": <str>, "cost_usd": <float|null>, "tests_count": <int>, "index": <int>, "total": <int>}
    {"type": "error",    "method": <str>, "file": <str>, "message": <str>, "index": <int|null>, "total": <int|null>}
    {"type": "summary",  "total": <int>,  "succeeded": <int>, "failed": <int>}
    {"type": "coverage_before", "success": <bool>, "coverage": <CoverageReport|absent>, "source_files_count": <int>, "test_files_count": <int>, "error": <str|absent>}
    {"type": "coverage_after",  "success": <bool>, "coverage": <CoverageReport|absent>, "source_files_count": <int>, "test_files_count": <int>, "error": <str|absent>}

Usage (matches early-cli env var naming):
    API_KEY=<key> py-agent generate-project --target-directory src/
"""

from __future__ import annotations

import argparse
import asyncio
import glob
import json
import logging
import os
import shutil
import sys
from typing import Any

from py_agent import PyAgent, PyAgentConfig
from py_agent.api.py_agent import TestableCompleteResult
from py_agent.infrastructure.coverage_runner import parse_coverage_report, run_targeted_coverage
from py_agent.models.testable import Testable
from py_agent.utils.git_diff import get_changed_line_ranges, ranges_overlap
from py_agent.utils.test_file_utils import compute_test_file_path, find_early_test_files


def _emit(event: dict[str, Any]) -> None:
    """Write one JSON event to stdout, flushed immediately."""
    print(json.dumps(event), flush=True)


def _is_test_file(path: str) -> bool:
    name = os.path.basename(path)
    if name.startswith("test_") or name.endswith("_test.py") or "_early" in name:
        return True
    parent = os.path.basename(os.path.dirname(path))
    return parent.startswith("test_early_") or parent.startswith("test_")


async def _run_generate_project(args: argparse.Namespace) -> int:
    config = PyAgentConfig(
        api_key=args.api_key,
        api_base_url=args.api_base_url,
        test_framework=args.test_framework,
        test_structure=args.test_structure,
        test_file_name=args.test_file_naming,
        concurrency=args.max_concurrency,
        workflow_run_id=args.workflow_run_id,
    )
    agent = PyAgent(config)
    await agent.init()

    target = os.path.abspath(args.target_directory)
    project_path = os.path.abspath(args.project_path or os.getcwd())

    # ── Phase 1: File discovery ──────────────────────────────────────────────
    if os.path.isfile(target) and target.endswith(".py"):
        py_files = [target] if not _is_test_file(target) else []
    else:
        py_files = [
            f
            for f in glob.glob(os.path.join(target, "**", "*.py"), recursive=True)
            if not _is_test_file(f)
        ]

    rel_target = os.path.relpath(target, project_path)
    _emit({"type": "scan", "files": len(py_files), "target": rel_target})

    # ── Phase 2: Testables extraction ────────────────────────────────────────
    testables_to_generate: list[dict[str, Any]] = []
    skipped_methods: list[dict[str, str]] = []
    all_source_files_with_testables: list[str] = []
    for file_path in py_files:
        try:
            all_testables = await agent.get_testables(file_path, project_path=project_path)
            testables = [t for t in all_testables if t.can_create_tests and t.type != "class"]
            classes_count = sum(1 for t in all_testables if t.type == "class")
            private_count = sum(
                1 for t in all_testables if not t.can_create_tests and t.type != "class"
            )
            rel_file = os.path.relpath(file_path, project_path)

            if testables:
                all_source_files_with_testables.append(file_path)

            # Filter out testables whose test file already exists on disk
            to_generate: list[Testable] = []
            skipped = 0
            for t in testables:
                cls = t.parent_name if t.type == "method" else None
                test_path = compute_test_file_path(
                    source_file_path=file_path,
                    project_path=project_path,
                    method_name=t.name,
                    test_structure=config.test_structure,
                    test_file_name=config.test_file_name,
                    class_name=cls,
                )
                if os.path.isfile(test_path):
                    skipped += 1
                    skipped_methods.append({"method": t.name, "file": rel_file})
                else:
                    to_generate.append(t)

            _emit(
                {
                    "type": "extract",
                    "file": rel_file,
                    "total_in_file": len(all_testables),
                    "functions": len(to_generate),
                    **({"skipped": skipped} if skipped > 0 else {}),
                    **({"classes_skipped": classes_count} if classes_count > 0 else {}),
                    **({"private_skipped": private_count} if private_count > 0 else {}),
                }
            )
            for t in to_generate:
                testables_to_generate.append(
                    {"file_path": file_path, "testable": t, "project_path": project_path}
                )
        except Exception as exc:
            _emit({"type": "error", "file": file_path, "method": "", "message": str(exc)})

    # Apply max-testables limit
    max_testables = args.max_testables
    if max_testables and len(testables_to_generate) > max_testables:
        testables_to_generate = testables_to_generate[:max_testables]

    # Pre-assign indices for progress tracking
    for idx, item in enumerate(testables_to_generate):
        item["_index"] = idx + 1

    total = len(testables_to_generate)
    file_count = len({item["file_path"] for item in testables_to_generate})
    _emit({"type": "start", "total": total, "files": file_count})

    # Emit per-method skip events so the caller can display them
    for skipped_item in skipped_methods:
        _emit(
            {
                "type": "skip",
                "method": skipped_item["method"],
                "file": skipped_item["file"],
            }
        )

    if total == 0:
        _emit({"type": "summary", "total": 0, "succeeded": 0, "failed": 0})
        _run_post_generation_coverage(
            all_source_files_with_testables, project_path, config.test_structure, set()
        )
        return 0

    # ── Phase 3: Test generation ─────────────────────────────────────────────
    succeeded = 0
    failed = 0
    generated_test_files: set[str] = set()

    def on_start(item: dict[str, Any]) -> None:
        testable: Testable = item["testable"]
        rel_file = os.path.relpath(item["file_path"], project_path)
        _emit(
            {
                "type": "progress",
                "method": testable.name,
                "file": rel_file,
                "index": item["_index"],
                "total": total,
            }
        )

    def on_complete(item: dict[str, Any], result: TestableCompleteResult | None) -> None:
        nonlocal succeeded, failed
        testable: Testable = item["testable"]
        rel_file = os.path.relpath(item["file_path"], project_path)
        if result is not None:
            succeeded += 1
            generated_test_files.add(result.test_file_path)
            _emit(
                {
                    "type": "done",
                    "method": testable.name,
                    "file": rel_file,
                    "test_file": os.path.relpath(result.test_file_path, project_path),
                    "request_id": result.request_id,
                    "cost_usd": result.cost_usd,
                    "tests_count": result.tests_count,
                    "index": item["_index"],
                    "total": total,
                }
            )
        else:
            failed += 1
            _emit(
                {
                    "type": "error",
                    "method": testable.name,
                    "file": rel_file,
                    "message": "generation failed",
                    "index": item["_index"],
                    "total": total,
                }
            )

    await agent.bulk_generate_tests(
        testables_to_generate,
        on_testable_complete=on_complete,
        on_testable_start=on_start,
    )

    _emit({"type": "summary", "total": total, "succeeded": succeeded, "failed": failed})

    _run_post_generation_coverage(
        all_source_files_with_testables,
        project_path,
        config.test_structure,
        generated_test_files,
    )

    return 0 if failed == 0 else 1


async def _run_generate_pr(args: argparse.Namespace) -> int:
    """Generate tests only for functions changed in the PR."""
    config = PyAgentConfig(
        api_key=args.api_key,
        api_base_url=args.api_base_url,
        test_framework=args.test_framework,
        test_structure=args.test_structure,
        test_file_name=args.test_file_naming,
        concurrency=args.max_concurrency,
        workflow_run_id=args.workflow_run_id,
    )
    agent = PyAgent(config)
    await agent.init()

    project_path = os.path.abspath(args.project_path or os.getcwd())
    base_ref = args.base_ref
    changed_files_raw = args.changed_files

    if not changed_files_raw:
        _emit(
            {"type": "error", "method": "", "message": "changed-files is required for generate-pr"}
        )
        return 1

    # ── Phase 1: Filter changed files to Python sources ──────────────────────
    # Changed files come from early-cli (GitHub/Bitbucket API), as relative paths.
    all_changed = [f.strip() for f in changed_files_raw.split("\n") if f.strip()]
    py_changed = [f for f in all_changed if f.endswith(".py") and not _is_test_file(f)]

    _emit(
        {
            "type": "scan",
            "files": len(py_changed),
            "target": f"PR ({len(all_changed)} changed files)",
        }
    )

    if not py_changed:
        _emit({"type": "summary", "total": 0, "succeeded": 0, "failed": 0})
        return 0

    # Get changed line ranges via local git diff (for changed-function detection)
    line_ranges = {}
    if base_ref:
        line_ranges = get_changed_line_ranges(base_ref, py_changed, cwd=project_path)

    # ── Phase 2: Extract testables and filter to changed functions ───────────
    changed_testables: list[dict[str, Any]] = []
    unchanged_testables: list[dict[str, Any]] = []
    skipped_methods: list[dict[str, str]] = []
    all_source_files_with_testables: list[str] = []

    for rel_file in py_changed:
        abs_path = os.path.join(project_path, rel_file)
        if not os.path.isfile(abs_path):
            continue

        try:
            all_testables = await agent.get_testables(abs_path, project_path=project_path)
            testables = [t for t in all_testables if t.can_create_tests and t.type != "class"]
            classes_count = sum(1 for t in all_testables if t.type == "class")
            private_count = sum(
                1 for t in all_testables if not t.can_create_tests and t.type != "class"
            )

            if testables:
                all_source_files_with_testables.append(abs_path)

            changed_in_file: list[Testable] = []
            unchanged_in_file: list[Testable] = []
            skipped = 0

            file_ranges = line_ranges.get(rel_file, [])

            for t in testables:
                # Skip if test file already exists
                cls = t.parent_name if t.type == "method" else None
                test_path = compute_test_file_path(
                    source_file_path=abs_path,
                    project_path=project_path,
                    method_name=t.name,
                    test_structure=config.test_structure,
                    test_file_name=config.test_file_name,
                    class_name=cls,
                )
                if os.path.isfile(test_path):
                    skipped += 1
                    skipped_methods.append({"method": t.name, "file": rel_file})
                    continue

                # Classify: changed vs unchanged based on line overlap
                # If no line ranges available (new file or git diff failed), treat all as changed
                if file_ranges and not ranges_overlap(
                    t.position.start.line,
                    t.position.end.line,
                    file_ranges,
                ):
                    unchanged_in_file.append(t)
                else:
                    changed_in_file.append(t)

            _emit(
                {
                    "type": "extract",
                    "file": rel_file,
                    "total_in_file": len(all_testables),
                    "functions": len(changed_in_file) + len(unchanged_in_file),
                    "changed": len(changed_in_file),
                    "unchanged": len(unchanged_in_file),
                    **({"skipped": skipped} if skipped > 0 else {}),
                    **({"classes_skipped": classes_count} if classes_count > 0 else {}),
                    **({"private_skipped": private_count} if private_count > 0 else {}),
                }
            )
            for t in changed_in_file:
                changed_testables.append(
                    {"file_path": abs_path, "testable": t, "project_path": project_path}
                )
            for t in unchanged_in_file:
                unchanged_testables.append(
                    {"file_path": abs_path, "testable": t, "project_path": project_path}
                )
        except Exception as exc:
            _emit({"type": "error", "file": rel_file, "method": "", "message": str(exc)})

    # Changed methods first (priority), then unchanged to fill remaining slots.
    # LimitFilter below will cap at max_testables, giving changed methods priority.
    testables_to_generate = changed_testables + unchanged_testables

    # Apply max-testables limit
    max_testables = args.max_testables
    if max_testables and len(testables_to_generate) > max_testables:
        testables_to_generate = testables_to_generate[:max_testables]

    # Pre-assign indices for progress tracking
    for idx, item in enumerate(testables_to_generate):
        item["_index"] = idx + 1

    total = len(testables_to_generate)
    file_count = len({item["file_path"] for item in testables_to_generate})
    _emit({"type": "start", "total": total, "files": file_count})

    for skipped_item in skipped_methods:
        _emit(
            {
                "type": "skip",
                "method": skipped_item["method"],
                "file": skipped_item["file"],
            }
        )

    if total == 0:
        _emit({"type": "summary", "total": 0, "succeeded": 0, "failed": 0})
        _run_post_generation_coverage(
            all_source_files_with_testables, project_path, config.test_structure, set()
        )
        return 0

    # ── Phase 3: Test generation ─────────────────────────────────────────────
    succeeded = 0
    failed = 0
    generated_test_files: set[str] = set()

    def on_start(item: dict[str, Any]) -> None:
        testable: Testable = item["testable"]
        rel_file = os.path.relpath(item["file_path"], project_path)
        _emit(
            {
                "type": "progress",
                "method": testable.name,
                "file": rel_file,
                "index": item["_index"],
                "total": total,
            }
        )

    def on_complete(item: dict[str, Any], result: TestableCompleteResult | None) -> None:
        nonlocal succeeded, failed
        testable: Testable = item["testable"]
        rel_file = os.path.relpath(item["file_path"], project_path)
        if result is not None:
            succeeded += 1
            generated_test_files.add(result.test_file_path)
            _emit(
                {
                    "type": "done",
                    "method": testable.name,
                    "file": rel_file,
                    "test_file": os.path.relpath(result.test_file_path, project_path),
                    "request_id": result.request_id,
                    "cost_usd": result.cost_usd,
                    "tests_count": result.tests_count,
                    "index": item["_index"],
                    "total": total,
                }
            )
        else:
            failed += 1
            _emit(
                {
                    "type": "error",
                    "method": testable.name,
                    "file": rel_file,
                    "message": "generation failed",
                    "index": item["_index"],
                    "total": total,
                }
            )

    await agent.bulk_generate_tests(
        testables_to_generate,
        on_testable_complete=on_complete,
        on_testable_start=on_start,
    )

    _emit({"type": "summary", "total": total, "succeeded": succeeded, "failed": failed})

    _run_post_generation_coverage(
        all_source_files_with_testables,
        project_path,
        config.test_structure,
        generated_test_files,
    )

    return 0 if failed == 0 else 1


def _run_post_generation_coverage(
    source_files: list[str],
    project_path: str,
    test_structure: str,
    generated_test_files: set[str],
) -> None:
    """Run coverage before/after for source files that had testables, using only early test files.

    "after" uses all early test files on disk for the targeted source files.
    "before" uses the same set minus the test files generated in this run.
    """
    if not source_files:
        return

    # Find all early test files on disk for those source files
    all_test_files: list[str] = []
    for src in source_files:
        all_test_files.extend(find_early_test_files(src, project_path, test_structure))

    source_count = len(source_files)

    if not all_test_files:
        _emit(
            {
                "type": "coverage",
                "success": False,
                "error": "No early test files found on disk",
                "source_files_count": source_count,
                "test_files_count": 0,
            }
        )
        return

    coverage_dir = os.path.join(project_path, ".early_coverage")

    # ── Coverage before: all early tests minus those generated this run ────
    before_test_files = [tf for tf in all_test_files if tf not in generated_test_files]
    before_result = _run_coverage_phase(
        "before",
        source_files,
        before_test_files,
        project_path,
        coverage_dir,
        source_count,
    )

    # ── Coverage after: skip if no new tests were generated (same data) ──
    if not generated_test_files:
        if before_result is not None:
            _emit({"type": "coverage_after", **before_result})
        else:
            _emit(
                {
                    "type": "coverage_after",
                    "success": True,
                    "coverage": {},
                    "source_files_count": source_count,
                    "test_files_count": len(before_test_files),
                }
            )
    else:
        _run_coverage_phase(
            "after",
            source_files,
            all_test_files,
            project_path,
            coverage_dir,
            source_count,
        )

    # Clean up coverage artifacts to avoid committing them
    shutil.rmtree(coverage_dir, ignore_errors=True)


def _run_coverage_phase(
    phase: str,
    source_files: list[str],
    test_files: list[str],
    project_path: str,
    coverage_dir: str,
    source_count: int,
) -> dict[str, Any] | None:
    """Run a single coverage phase (before or after) and emit the event.

    Returns the event payload (without ``type``) on success so callers can
    reuse the data, or ``None`` on failure / empty.
    """
    test_count = len(test_files)

    if not test_files:
        event_data: dict[str, Any] = {
            "success": True,
            "coverage": {},
            "source_files_count": source_count,
            "test_files_count": 0,
        }
        _emit({"type": f"coverage_{phase}", **event_data})
        return event_data

    coverage_report_path = os.path.join(coverage_dir, f"coverage_{phase}.json")

    result = run_targeted_coverage(
        project_path=project_path,
        coverage_report_path=coverage_report_path,
        source_files=source_files,
        test_files=test_files,
    )

    if not result.get("success"):
        _emit(
            {
                "type": f"coverage_{phase}",
                "success": False,
                "error": result.get("error", "Coverage run failed"),
                "source_files_count": source_count,
                "test_files_count": test_count,
            }
        )
        return None

    try:
        coverage_tree = parse_coverage_report(coverage_report_path, project_path, source_files)
        event_data = {
            "success": True,
            "coverage": coverage_tree,
            "source_files_count": source_count,
            "test_files_count": test_count,
        }
        _emit({"type": f"coverage_{phase}", **event_data})
        return event_data
    except Exception as exc:
        _emit(
            {
                "type": f"coverage_{phase}",
                "success": False,
                "error": f"Failed to parse coverage report: {exc}",
                "source_files_count": source_count,
                "test_files_count": test_count,
            }
        )
        return None


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="py-agent",
        description="EarlyAI py-agent — generates tests for Python projects",
    )

    # Auth
    parser.add_argument(
        "--api-key",
        default=os.environ.get("API_KEY"),
        help="EarlyAI API key (env: API_KEY)",
    )
    parser.add_argument(
        "--api-base-url",
        default=os.environ.get("API_BASE_URL", "https://api.startearly.ai"),
        help="Backend base URL (env: API_BASE_URL)",
    )

    # Paths
    parser.add_argument(
        "--target-directory",
        default=os.environ.get("TARGET_DIRECTORY", "."),
        help="Directory to generate tests for (env: TARGET_DIRECTORY)",
    )
    parser.add_argument(
        "--project-path",
        default=os.environ.get("PROJECT_PATH"),
        help="Project root path (env: PROJECT_PATH, default: cwd)",
    )

    # Test config
    parser.add_argument(
        "--test-framework",
        default=os.environ.get("TEST_FRAMEWORK", "pytest"),
        choices=["pytest", "unittest"],
        help="Test framework (env: TEST_FRAMEWORK)",
    )
    parser.add_argument(
        "--test-structure",
        default=os.environ.get("TEST_STRUCTURE", "siblingFolder"),
        choices=["siblingFolder", "rootFolder"],
        help="Test file structure (env: TEST_STRUCTURE)",
    )
    parser.add_argument(
        "--test-file-naming",
        default=os.environ.get("TEST_FILE_NAMING", "camelCase"),
        choices=["camelCase", "kebabCase"],
        help="File naming convention (env: TEST_FILE_NAMING)",
    )
    parser.add_argument(
        "--max-concurrency",
        type=int,
        default=int(os.environ.get("MAX_CONCURRENCY", "3")),
        help="Concurrent workers (env: MAX_CONCURRENCY)",
    )

    parser.add_argument(
        "--max-testables",
        type=int,
        default=int(os.environ.get("MAX_TESTABLES", "0")) or None,
        help="Max number of methods to generate tests for (env: MAX_TESTABLES, 0=unlimited)",
    )

    parser.add_argument(
        "--workflow-run-id",
        default=os.environ.get("WORKFLOW_RUN_ID"),
        help="Workflow run ID for linking operations (env: WORKFLOW_RUN_ID)",
    )

    # PR-specific options
    parser.add_argument(
        "--base-ref",
        default=os.environ.get("BASE_REF"),
        help="Base branch/ref for PR diff (env: BASE_REF)",
    )
    parser.add_argument(
        "--changed-files",
        default=os.environ.get("CHANGED_FILES"),
        help="Newline-separated list of changed file paths from SCM provider (env: CHANGED_FILES)",
    )

    parser.add_argument(
        "command",
        nargs="?",
        default="generate-project",
        choices=["generate-project", "generate-pr"],
        help="Command to run (default: generate-project)",
    )

    return parser


def main() -> None:
    log_level = os.environ.get("LOG_LEVEL", "INFO").upper()
    logging.basicConfig(
        level=getattr(logging, log_level, logging.INFO),
        format="%(levelname)s %(name)s: %(message)s",
        stream=sys.stderr,
    )

    parser = _build_parser()
    args = parser.parse_args()

    logging.getLogger(__name__).debug(
        f"WORKFLOW_RUN_ID env={os.environ.get('WORKFLOW_RUN_ID')!r} arg={getattr(args, 'workflow_run_id', None)!r}"
    )

    # Default to generate-project when no subcommand given
    if args.command is None:
        args.command = "generate-project"

    if not args.api_key:
        _emit({"type": "error", "message": "API key required: set API_KEY env var or --api-key"})
        sys.exit(1)

    if args.command == "generate-project":
        sys.exit(asyncio.run(_run_generate_project(args)))
    elif args.command == "generate-pr":
        sys.exit(asyncio.run(_run_generate_pr(args)))
    else:
        parser.print_help()
        sys.exit(1)
