# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
"""
Unittestcasegenerator Environment Implementation.

An RL environment where an LLM agent writes JUnit 5 unit tests
for given Java classes. Tasks range from simple pure functions
to stateful classes requiring exception testing.
"""

import os
import re
import subprocess
import shutil
import tempfile
import textwrap
import json
from typing import Any, Optional
from uuid import uuid4

# from openenv.core.env_server.interfaces import Environment
# from openenv.core.env_server.types import State


try:
    from openenv.core.env_server.interfaces import Environment
    from openenv.core.env_server.types import State
except ImportError:
    # Stub classes so the module can be imported without openenv installed
    class Environment:
        pass

    class State:
        def __init__(self, episode_id=None, step_count=0):
            self.episode_id = episode_id
            self.step_count = step_count
try:
    from ..models import UnittestcasegeneratorAction, UnittestcasegeneratorObservation
except ImportError:
    from models import UnittestcasegeneratorAction, UnittestcasegeneratorObservation


# ─────────────────────────────────────────────
# TASKS
# ─────────────────────────────────────────────

TASKS = {
    "easy": {
        "source_code": textwrap.dedent("""
            public class Calculator {
                public int add(int a, int b) {
                    return a + b;
                }
                public boolean isPalindrome(String s) {
                    String clean = s.toLowerCase().replace(" ", "");
                    String reversed = new StringBuilder(clean).reverse().toString();
                    return clean.equals(reversed);
                }
                public double celsiusToFahrenheit(double c) {
                    return c * 9.0 / 5.0 + 32;
                }
            }
        """).strip(),
        "source_class": "Calculator",
        "test_class": "CalculatorTest",
        "task_hint": (
            "Write JUnit 5 tests for Calculator class. "
            "1) Class name MUST be CalculatorTest. "
            "2) Use @Test annotation on every test method. "
            "3) Import: import org.junit.jupiter.api.Test; "
            "4) Import: import static org.junit.jupiter.api.Assertions.*; "
            "5) Write at least 6 test methods. "
            "6) Test add(), isPalindrome(), celsiusToFahrenheit(). "
            "Reply with ONLY the Java code."
        ),
        "expected_min_tests": 6,
        "requires_edge_cases": False,
        "requires_exception": False,
    },

    "medium": {
        "source_code": textwrap.dedent("""
            public class SafeMath {
                public double safeDivide(double a, double b) {
                    if (b == 0)
                        throw new IllegalArgumentException("Cannot divide by zero");
                    return a / b;
                }
                public int getFirstElement(int[] arr) {
                    if (arr.length == 0)
                        throw new IndexOutOfBoundsException("Array is empty");
                    return arr[0];
                }
                public int parsePositiveInt(String s) {
                    int val = Integer.parseInt(s);
                    if (val <= 0)
                        throw new IllegalArgumentException("Must be positive");
                    return val;
                }
            }
        """).strip(),
        "source_class": "SafeMath",
        "test_class": "SafeMathTest",
        "task_hint": (
            "Write JUnit 5 tests for SafeMath class. "
            "1) Class name MUST be SafeMathTest. "
            "2) Use @Test annotation on every test method. "
            "3) Import: import org.junit.jupiter.api.Test; "
            "4) Import: import static org.junit.jupiter.api.Assertions.*; "
            "5) Use assertThrows() for exception testing. "
            "6) Test edge cases: empty array, zero, negative. "
            "7) Write at least 8 test methods. "
            "Reply with ONLY the Java code."
        ),
        "expected_min_tests": 8,
        "requires_edge_cases": True,
        "requires_exception": True,
    },

    "hard": {
        "source_code": textwrap.dedent("""
            import java.util.ArrayList;
            import java.util.List;
            public class BankAccount {
                private String owner;
                private double balance;
                private List<String> transactions;
                public BankAccount(String owner, double balance) {
                    this.owner = owner;
                    this.balance = balance;
                    this.transactions = new ArrayList<>();
                }
                public double deposit(double amount) {
                    if (amount <= 0)
                        throw new IllegalArgumentException("Deposit must be positive");
                    this.balance += amount;
                    this.transactions.add("deposit:" + amount);
                    return this.balance;
                }
                public double withdraw(double amount) {
                    if (amount <= 0)
                        throw new IllegalArgumentException("Withdrawal must be positive");
                    if (amount > this.balance)
                        throw new IllegalArgumentException("Insufficient funds");
                    this.balance -= amount;
                    this.transactions.add("withdraw:" + amount);
                    return this.balance;
                }
                public int getTransactionCount() { return this.transactions.size(); }
                public double getBalance() { return this.balance; }
                public String getOwner() { return this.owner; }
            }
        """).strip(),
        "source_class": "BankAccount",
        "test_class": "BankAccountTest",
        "task_hint": (
            "Write JUnit 5 tests for BankAccount class. "
            "1) Class name MUST be BankAccountTest. "
            "2) Use @Test annotation on every test method. "
            "3) Import: import org.junit.jupiter.api.Test; "
            "4) Import: import static org.junit.jupiter.api.Assertions.*; "
            "5) Use assertThrows() for exception testing. "
            "6) Test deposit, withdraw, getBalance, getTransactionCount, getOwner. "
            "7) Test exceptions: negative deposit, withdraw, insufficient funds. "
            "8) Write at least 10 test methods. "
            "Reply with ONLY the Java code."
        ),
        "expected_min_tests": 10,
        "requires_edge_cases": True,
        "requires_exception": True,
    },
}

DIFFICULTIES = ("easy", "medium", "hard")


# ─────────────────────────────────────────────
# JAVA DETECTION
# ─────────────────────────────────────────────

def find_java():
    for cmd in ["java", "/usr/bin/java", "/usr/local/bin/java"]:
        try:
            result = subprocess.run([cmd, "-version"], capture_output=True, text=True)
            if result.returncode == 0:
                print(json.dumps({"event": "debug", "type": "java_found", "path": cmd}), flush=True)
                return cmd
        except FileNotFoundError:
            continue
    return None


def find_javac():
    for cmd in ["javac", "/usr/bin/javac", "/usr/local/bin/javac"]:
        try:
            result = subprocess.run([cmd, "-version"], capture_output=True, text=True)
            if result.returncode == 0:
                print(json.dumps({"event": "debug", "type": "javac_found", "path": cmd}), flush=True)
                return cmd
        except FileNotFoundError:
            continue
    return None


def find_junit_jar():
    env_path = os.getenv("JUNIT_JAR")
    if env_path and os.path.exists(env_path):
        print(json.dumps({"event": "debug", "type": "junit_found", "path": env_path}), flush=True)
        return env_path

    candidates = [
        "/app/junit-platform-console-standalone.jar",
        "/app/env/junit-platform-console-standalone.jar",
        os.path.join(os.path.dirname(__file__), "..", "junit-platform-console-standalone.jar"),
        os.path.join(os.getcwd(), "junit-platform-console-standalone.jar"),
        os.path.join(os.getcwd(), "libs", "junit-platform-console-standalone.jar"),
        # "/Users/vidhikoul/Desktop/UnitTestCaseGenerator/tcgenerator/junit-platform-console-standalone.jar",
    ]

    for path in candidates:
        abs_path = os.path.abspath(path)
        if os.path.exists(abs_path):
            print(json.dumps({"event": "debug", "type": "junit_found", "path": abs_path}), flush=True)
            return abs_path

    print(json.dumps({"event": "debug", "type": "junit_not_found", "tried": candidates}), flush=True)
    return None


# ─────────────────────────────────────────────
# JUNIT SANDBOX
# ─────────────────────────────────────────────

def run_junit_tests(source_code: str, test_code: str,
                    source_class: str, test_class: str,
                    timeout: int = 30):

    java  = find_java()
    javac = find_javac()
    junit = find_junit_jar()

    if not java or not javac:
        return 0, 0, 0, "Java not found!"
    if not junit:
        return 0, 0, 0, "JUnit jar not found!"

    tmpdir = tempfile.mkdtemp()

    try:
        source_file = os.path.join(tmpdir, f"{source_class}.java")
        test_file   = os.path.join(tmpdir, f"{test_class}.java")

        with open(source_file, "w") as f:
            f.write(source_code)
        with open(test_file, "w") as f:
            f.write(test_code)

        compile_result = subprocess.run(
            [javac, "-cp", junit, source_file, test_file],
            capture_output=True, text=True,
            cwd=tmpdir, timeout=30,
        )

        if compile_result.returncode != 0:
            err = compile_result.stderr[:500]
            print(json.dumps({"event": "debug", "type": "compile_error", "message": err[:200]}), flush=True)
            return 0, 0, 0, f"Compile error: {err}"

        run_result = subprocess.run(
            [java, "-jar", junit, "-cp", tmpdir,
             "--select-class", test_class, "--details", "summary"],
            capture_output=True, text=True,
            cwd=tmpdir, timeout=timeout,
        )

        output = run_result.stdout + run_result.stderr
        print(json.dumps({"event": "debug", "type": "junit_output", "message": output[:300]}), flush=True)

        passed, failed = parse_junit_output(output)
        total = passed + failed
        error = None if passed > 0 else output[-400:]

        return passed, failed, total, error

    except subprocess.TimeoutExpired:
        return 0, 0, 0, "Timeout"
    except Exception as e:
        return 0, 0, 0, str(e)
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


def parse_junit_output(output: str):
    passed = failed = 0

    p = re.findall(r'(\d+)\s+tests?\s+successful', output, re.IGNORECASE)
    if p:
        passed = int(p[0])

    f = re.findall(r'(\d+)\s+tests?\s+failed', output, re.IGNORECASE)
    if f:
        failed = int(f[0])

    if passed == 0 and failed == 0:
        passed = len(re.findall(r'\[\s*OK\s*\]', output))
        failed = len(re.findall(r'\[\s*FAILED\s*\]', output))

    return passed, failed


def compute_reward(passed, total, test_code, task_cfg):
    eps = 0.001

    if total == 0:
        return eps

    base = (passed / total) * 0.7
    quantity_bonus = 0.1 if total >= task_cfg["expected_min_tests"] else 0.0

    edge_bonus = 0.0
    if task_cfg["requires_edge_cases"]:
        edge_keywords = ["empty", "zero", "null", "negative", "0", "[]", '""']
        if any(kw in test_code.lower() for kw in edge_keywords):
            edge_bonus = 0.1

    exception_bonus = 0.0
    if task_cfg["requires_exception"]:
        if "assertthrows" in test_code.lower():
            exception_bonus = 0.1

    reward = base + quantity_bonus + edge_bonus + exception_bonus

    reward = max(eps, min(reward, 1 - eps))

    return reward
# ─────────────────────────────────────────────
# ENVIRONMENT CLASS
# ─────────────────────────────────────────────

class UnittestcasegeneratorEnvironment(Environment):
    """
    JUnit Test Case Generator Environment.

    An RL environment where an LLM agent learns to write JUnit 5 unit
    tests for Java classes. The agent receives Java source code and must
    produce test code that compiles and passes.

    Difficulty levels:
      easy:   Simple pure functions — Calculator class (6+ tests)
      medium: Exception handling   — SafeMath class (8+ tests)
      hard:   Stateful class       — BankAccount class (10+ tests)

    Episode flow:
      1. reset(difficulty="easy"|"medium"|"hard")
      2. agent reads source_code + task_hint from observation
      3. agent calls step(action) with JUnit test code
      4. environment compiles + runs tests → returns reward (0.0–1.0)
      5. repeat up to 6 steps or until reward >= 0.95
    """

    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self):
        """Initialize the UnitTestCaseGenerator environment."""
        self._state         = State(episode_id=str(uuid4()), step_count=0)
        self._episode_count = 0
        self._step_count    = 0
        self._max_steps     = 6
        self._best_reward   = 0.0
        self._difficulty    = "easy"
        self._task_cfg      = TASKS["easy"]

    def reset(
        self,
        difficulty: Optional[str] = None,
        episode_id: Optional[str] = None,
        **kwargs: Any,
    ) -> UnittestcasegeneratorObservation:
        """
        Reset the environment with a new task.

        Args:
            difficulty: "easy" | "medium" | "hard"
                        Auto-cycles easy→medium→hard if not provided.
            episode_id: Optional custom episode ID.

        Returns:
            UnittestcasegeneratorObservation with source_code and task_hint.
        """
        if difficulty not in DIFFICULTIES:
            difficulty = DIFFICULTIES[self._episode_count % 3]

        self._episode_count += 1
        self._difficulty    = difficulty
        self._task_cfg      = TASKS[difficulty]
        self._step_count    = 0
        self._best_reward   = 0.0
        self._state         = State(
            episode_id=episode_id or str(uuid4()),
            step_count=0
        )

        print(json.dumps({
            "event": "debug",
            "type": "reset",
            "difficulty": difficulty,
            "source_class": self._task_cfg["source_class"],
        }), flush=True)

        return UnittestcasegeneratorObservation(
            source_code = self._task_cfg["source_code"],
            task_hint   = self._task_cfg["task_hint"],
            passed=0, failed=0, total=0,
            error=None, reward=0.0, done=False,
            metadata={
                "difficulty": difficulty,
                "episode_id": self._state.episode_id,
                "message": (
                    f"New {difficulty} task loaded. "
                    f"Write JUnit 5 tests for {self._task_cfg['source_class']} class."
                ),
            }
        )

    def step(self, action: UnittestcasegeneratorAction) -> UnittestcasegeneratorObservation:
        """
        Execute one step — compile and run the agent's JUnit test code.

        Args:
            action: UnittestcasegeneratorAction with test_code field.

        Returns:
            UnittestcasegeneratorObservation with passed/failed counts and reward.
        """
        self._state.step_count += 1
        self._step_count += 1
        done = self._step_count >= self._max_steps

        passed, failed, total, error = run_junit_tests(
            source_code  = self._task_cfg["source_code"],
            test_code    = action.test_code,
            source_class = self._task_cfg["source_class"],
            test_class   = self._task_cfg["test_class"],
        )

        reward = compute_reward(passed, total, action.test_code, self._task_cfg)
        self._best_reward = max(self._best_reward, reward)

        if reward >= 0.95:
            done = True

        print(json.dumps({
            "event": "debug",
            "type": "step_result",
            "passed": passed,
            "failed": failed,
            "total": total,
            "reward": reward,
            "done": done,
        }), flush=True)

        return UnittestcasegeneratorObservation(
            source_code = self._task_cfg["source_code"],
            task_hint   = self._task_cfg["task_hint"],
            passed=passed, failed=failed, total=total,
            error=error, reward=reward, done=done,
        )

    @property
    def state(self) -> State:
        """Get current episode state."""
        return self._state