"""Generate Harbor-format task directories from the resolved catalog.

Usage:
    python rl_env/task_generator.py \
        --resolved-catalog resolved_catalog_full.json \
        --output-dir codeflash-perf-tasks \
        --inference-repo /path/to/inference
"""

from __future__ import annotations

import argparse
import json
import logging
import re
import shutil
import subprocess
import textwrap
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
log = logging.getLogger(__name__)


def sanitize_task_name(function_name: str, file_path: str) -> str:
    """Create a filesystem-safe task directory name."""
    # Use function name + last path component for uniqueness
    parent = Path(file_path).stem  # e.g., "v1" from ".../v1.py"
    grandparent = Path(file_path).parent.name  # e.g., "detection_event_log"
    name = f"{grandparent}-{function_name}".replace(".", "-")
    # Sanitize
    name = re.sub(r"[^a-zA-Z0-9_-]", "-", name)
    name = re.sub(r"-+", "-", name).strip("-").lower()
    return name[:80]  # cap length


def find_function_line_number(file_content: str, function_name: str) -> int | None:
    """Find the line number where a function is defined."""
    short_name = function_name.rsplit(".", maxsplit=1)[-1]
    for i, line in enumerate(file_content.splitlines(), 1):
        stripped = line.lstrip()
        if stripped.startswith(f"def {short_name}(") or stripped.startswith(f"async def {short_name}("):
            return i
    return None


def generate_task_toml(entry: dict) -> str:
    func_fqn = f"{entry['file_path'].replace('/', '.').removesuffix('.py')}.{entry['function_name']}"
    tags = ["optimization", "python"]
    if "workflows" in entry["file_path"]:
        tags.append("workflows")
    if "models" in entry["file_path"]:
        tags.append("ml-models")
    if "visualizations" in entry["file_path"]:
        tags.append("visualization")

    return textwrap.dedent(f"""\
        version = "1.0"

        [metadata]
        author_name = "Codeflash"
        author_email = "saurabh@codeflash.ai"
        category = "code-performance"
        language = "python"
        repo = "roboflow/inference"
        function = "{func_fqn}"
        tags = {json.dumps(tags)}

        [environment]
        cpus = 4
        memory_mb = 8192
        storage_mb = 20480
        gpus = 0
        allow_internet = false

        [agent]
        timeout_sec = 1200

        [verifier]
        timeout_sec = 600
    """)


def generate_instruction_md(entry: dict, line_number: int | None) -> str:
    func_name = entry["function_name"]
    file_path = entry["file_path"]
    line_str = f" (line {line_number})" if line_number else ""

    return textwrap.dedent(f"""\
        # Task: Optimize a Python function for performance

        You are working in the repository `roboflow/inference` (already cloned at `/workspace/inference`).

        **Target function:** `{func_name}` in `{file_path}`{line_str}

        Your goal is to **rewrite this function so it runs faster** while preserving
        identical behavior. The function must:
        - Return the same values for all inputs
        - Produce the same side effects (mutations, file writes, prints, etc.)
        - Handle the same edge cases and raise the same exceptions

        Edit the function **in place** in the repository. You may also:
        - Add new imports at the top of the file
        - Add helper functions in the same file
        - Modify the function signature's implementation but NOT its interface

        You have access to the full repository for context. You may inspect any file,
        run the existing test suite, or write your own tests. However, you will be
        graded on a held-out evaluation suite that you cannot see.

        When you are done, create a git commit with your changes:
        ```
        git commit -am "Optimized {func_name}"
        ```
    """)


def generate_solve_sh(entry: dict, task_dir: Path) -> str:
    """Generate solve.sh based on the resolution method."""
    method = entry["resolution_method"]

    if method == "git_parent" and entry.get("patch"):
        # We have a real git patch — use git apply
        return textwrap.dedent("""\
            #!/bin/bash
            set -e
            cd /workspace/inference
            git apply /solution/optimization.patch
            git commit -am "Apply known-good optimization"
        """)
    else:
        # Use the Python replacement script
        return textwrap.dedent("""\
            #!/bin/bash
            set -e
            cd /workspace/inference
            python /solution/apply_optimization.py
            git commit -am "Apply known-good optimization"
        """)


APPLY_TEMPLATE_PATH = Path(__file__).parent / "apply_optimization_template.py"


def generate_apply_config(entry: dict) -> str:
    """Generate config.json for the apply_optimization_template.py script."""
    short_name = entry["function_name"].rsplit(".", maxsplit=1)[-1]
    return json.dumps(
        {"target_file": entry["file_path"], "function_name": short_name},
        indent=2,
    )


def generate_test_sh() -> str:
    """Generate test.sh placeholder — will be replaced by evaluator in Phase 2."""
    return textwrap.dedent("""\
        #!/bin/bash
        set -e
        cd /workspace/inference

        # Placeholder: Codeflash evaluation harness (Phase 2)
        # For now, just verify the code is syntactically valid and tests pass
        python -c "
        import ast, sys
        # Verify all .py files that were modified can be parsed
        import subprocess
        result = subprocess.run(['git', 'diff', '--name-only', 'HEAD~1'], capture_output=True, text=True)
        for f in result.stdout.strip().split('\\n'):
            if f.endswith('.py'):
                try:
                    ast.parse(open(f).read())
                except SyntaxError as e:
                    print(f'SyntaxError in {f}: {e}')
                    sys.exit(1)
        print('All modified files pass syntax check')
        "

        # Write a basic reward.json
        mkdir -p /logs/verifier
        echo '{"correct": 1.0, "speedup": 0.0, "tests_passed": 0, "tests_total": 0}' > /logs/verifier/reward.json
        echo "Evaluation placeholder complete"
    """)


def generate_base_dockerfile() -> str:
    return textwrap.dedent("""\
        FROM python:3.12-slim

        RUN apt-get update && apt-get install -y --no-install-recommends \\
            git build-essential libgl1 libglib2.0-0 \\
            && rm -rf /var/lib/apt/lists/*

        # Clone inference repo
        RUN git clone https://github.com/roboflow/inference.git /workspace/inference
        WORKDIR /workspace/inference

        # Install inference dependencies (core only, no GPU)
        RUN pip install --no-cache-dir -e ".[test]" || \\
            pip install --no-cache-dir -e . 2>/dev/null || true

        # Install codeflash evaluator (Phase 2 — placeholder)
        # COPY codeflash-evaluator/ /opt/codeflash-eval/
        # RUN pip install -e /opt/codeflash-eval/
    """)


def generate_task_dockerfile(entry: dict) -> str:
    commit = entry["pre_optimization_commit"]
    if commit == "HEAD":
        checkout_cmd = "# Using HEAD — no checkout needed"
        reinstall_cmd = ""
    else:
        checkout_cmd = f"RUN cd /workspace/inference && git checkout {commit}"
        # Re-install inference at this commit to get correct dependency versions
        # (inference_models, pydantic, etc. may differ between commits)
        reinstall_cmd = "RUN pip install --no-cache-dir -e . 2>/dev/null || true"

    return textwrap.dedent(f"""\
        FROM codeflash-inference-base:latest

        # Checkout to pre-optimization state
        {checkout_cmd}

        # Re-install to get correct dependency versions for this commit
        {reinstall_cmd}

        WORKDIR /workspace/inference
    """)


def generate_task(entry: dict, output_dir: Path, inference_repo: Path) -> Path | None:
    """Generate a complete Harbor task directory for one optimization."""
    task_name = sanitize_task_name(entry["function_name"], entry["file_path"])
    task_dir = output_dir / task_name

    # Skip if already exists
    if task_dir.exists():
        log.debug(f"  Skipping {task_name} — already exists")
        return task_dir

    # Find line number in original file
    line_number = find_function_line_number(entry["original_file_content"], entry["function_name"])

    # Create directory structure
    task_dir.mkdir(parents=True)
    (task_dir / "environment").mkdir()
    (task_dir / "solution").mkdir()
    (task_dir / "tests").mkdir()
    (task_dir / "tests" / "codeflash_eval").mkdir()

    # task.toml
    (task_dir / "task.toml").write_text(generate_task_toml(entry), encoding="utf-8")

    # instruction.md
    (task_dir / "instruction.md").write_text(generate_instruction_md(entry, line_number), encoding="utf-8")

    # environment/Dockerfile
    (task_dir / "environment" / "Dockerfile").write_text(generate_task_dockerfile(entry), encoding="utf-8")

    # solution/solve.sh
    solve_sh = generate_solve_sh(entry, task_dir)
    (task_dir / "solution" / "solve.sh").write_text(solve_sh, encoding="utf-8")

    # solution/optimization.patch (for git_parent method)
    if entry.get("patch"):
        (task_dir / "solution" / "optimization.patch").write_text(entry["patch"], encoding="utf-8")

    # solution/apply_optimization.py + config (for db methods)
    if entry["resolution_method"] != "git_parent":
        shutil.copy2(APPLY_TEMPLATE_PATH, task_dir / "solution" / "apply_optimization.py")
        (task_dir / "solution" / "config.json").write_text(generate_apply_config(entry), encoding="utf-8")
        # solution/optimized_code.py — raw markdown with file path fences for multi-file support
        raw_md = entry.get("optimized_raw_markdown") or entry["optimized_function_code"]
        (task_dir / "solution" / "optimized_code.py").write_text(raw_md, encoding="utf-8")

    # tests/test.sh
    (task_dir / "tests" / "test.sh").write_text(generate_test_sh(), encoding="utf-8")

    # tests/codeflash_eval/original.py — full original file content for baseline
    # Using full file (not just function) so libcst can parse it for replacement
    (task_dir / "tests" / "codeflash_eval" / "original.py").write_text(
        entry["original_file_content"], encoding="utf-8"
    )

    return task_dir


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate Harbor task directories")
    parser.add_argument(
        "--resolved-catalog",
        type=Path,
        default=Path("resolved_catalog_full.json"),
        help="Input from commit_resolver.py",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("codeflash-perf-tasks"),
        help="Output directory for Harbor tasks",
    )
    parser.add_argument(
        "--inference-repo",
        type=Path,
        default=Path("/Users/saurabh/Dropbox/hacks/inference"),
        help="Path to local inference repo clone",
    )
    parser.add_argument("--limit", type=int, default=0, help="Limit tasks to generate (0 = all)")
    args = parser.parse_args()

    catalog = json.loads(args.resolved_catalog.read_text(encoding="utf-8"))
    log.info(f"Loaded {len(catalog)} resolved entries")

    if args.limit > 0:
        catalog = catalog[: args.limit]

    # Create output directory
    args.output_dir.mkdir(parents=True, exist_ok=True)

    # Generate base image Dockerfile
    base_dir = args.output_dir / "base-image"
    base_dir.mkdir(exist_ok=True)
    (base_dir / "Dockerfile").write_text(generate_base_dockerfile(), encoding="utf-8")

    # Generate tasks
    generated = 0
    failed = 0
    for i, entry in enumerate(catalog):
        task_dir = generate_task(entry, args.output_dir, args.inference_repo)
        if task_dir:
            generated += 1
        else:
            failed += 1

        if (i + 1) % 50 == 0:
            log.info(f"  Generated {i + 1}/{len(catalog)}...")

    log.info(f"Generated {generated} tasks, {failed} failed")
    log.info(f"Output: {args.output_dir}")

    # List a few task directories
    task_dirs = sorted([d for d in args.output_dir.iterdir() if d.is_dir() and d.name != "base-image"])
    log.info(f"Task directories: {len(task_dirs)}")
    for d in task_dirs[:10]:
        files = list(d.rglob("*"))
        log.info(f"  {d.name}/ ({len(files)} files)")


if __name__ == "__main__":
    main()
