"""
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>}

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 sys
from typing import Any

from py_agent import PyAgent, PyAgentConfig
from py_agent.api.py_agent import TestableCompleteResult
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


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)
    return name.startswith("test_") or name.endswith("_test.py") or "_early" in name


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]] = []
    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]
            rel_file = os.path.relpath(file_path, project_path)

            # Filter out testables whose test file already exists on disk
            to_generate: list[Testable] = []
            skipped = 0
            for t in testables:
                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,
                )
                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 {}),
                }
            )
            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})
        return 0

    # ── Phase 3: Test generation ─────────────────────────────────────────────
    succeeded = 0
    failed = 0

    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
            _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})
    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]] = []

    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]

            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
                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,
                )
                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 {}),
                }
            )
            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})
        return 0

    # ── Phase 3: Test generation ─────────────────────────────────────────────
    succeeded = 0
    failed = 0

    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
            _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})
    return 0 if failed == 0 else 1


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)
