#!/usr/bin/env python3
# Shebang line: specifies the interpreter to be used for executing the script (Python 3).
import os
# Imports the os module for interacting with the operating system, such as file paths and environment variables.
import sys
# Imports the sys module, providing access to system-specific parameters and functions, like sys.exit.
import json
# Imports the json module for working with JSON data (encoding and decoding).
import subprocess
# Imports the subprocess module for running new processes (shell commands).
from datetime import datetime
# Imports the datetime class from the datetime module, used for generating timestamps.


def sh(cmd: str) -> subprocess.CompletedProcess:
    # Function to execute a shell command.
    # cmd: The shell command string to execute.
    # -> subprocess.CompletedProcess: Returns the result of the executed process.
    return subprocess.run(
        cmd,                     # The command to execute.
        shell=True,              # Execute the command through the shell (allows shell features like piping).
        stdout=subprocess.PIPE,  # Captures standard output.
        stderr=subprocess.STDOUT, # Directs standard error to standard output (so both are captured in stdout).
        text=True                # Decodes stdout and stderr as text using the default encoding.
    )


def list_tree(root: str, max_depth: int = 2, max_entries: int = 500) -> str:
    # Function to generate a simplified file tree structure string for a given root directory.
    # root: The starting directory path.
    # max_depth: The maximum directory depth to traverse (default 2).
    # max_entries: The maximum total number of file entries to include (default 500).
    # -> str: Returns the file tree as a single string, with lines separated by newlines.
    out_lines = []
    # List to store the lines of the file tree output.
    count = 0
    # Counter for the number of files added to the output.
    for cur_root, dirs, files in os.walk(root):
        # os.walk iterates through the directory tree rooted at 'root'.
        # cur_root: The current directory path.
        # dirs: A list of subdirectory names in cur_root.
        # files: A list of file names in cur_root.
        
        # Calculate the current depth relative to the initial root.
        depth = cur_root[len(root):].count(os.sep)
        # Calculate indentation based on depth.
        indent = "  " * depth
        
        # Add the current directory name to the output with indentation and a trailing slash.
        out_lines.append(f"{indent}{os.path.basename(cur_root)}/")
        
        # Process and list files in the current directory, limiting to the first 200 files in this directory.
        for f in sorted(files)[:200]:
            # Check if the maximum total entry count has been reached.
            if count >= max_entries:
                # If max entries reached, append a truncation message and return the output.
                out_lines.append("  ... (truncated)")
                return "\n".join(out_lines)
            
            # Add the file name to the output with indentation.
            out_lines.append(f"{indent}  {f}")
            # Increment the file count.
            count += 1
            
        # Prune the directory list if the maximum depth is reached.
        if depth >= max_depth:
            # Setting dirs[:] = [] prevents os.walk from descending into subdirectories of the current directory.
            dirs[:] = []
            
    # Join all collected lines into a single string and return it.
    return "\n".join(out_lines)


def main() -> int:
    # Main function to execute the LLM setup helper logic.
    # -> int: Returns an exit code (0 for success, non-zero for error).
    
    # Get the Anthropic API key from the environment variable.
    api_key = os.getenv("ANTHROPIC_API_KEY", "").strip()
    
    # Check if the API key is available. If not, log a message and exit gracefully (assuming no LLM interaction is possible).
    if not api_key:
        print("ANTHROPIC_API_KEY not set; skipping LLM setup helper.")
        return 0

    # Get work repository path from environment, defaulting to a specific path.
    work_repo = os.environ.get("WORK_REPO", "/home/ubuntu/work/repo")
    # Get the root directory for tests from environment variables, preferring LLM_TESTS_ROOT.
    tests_root = os.environ.get("LLM_TESTS_ROOT") or os.environ.get("TESTS_ROOT") or ""
    # Get the pytest command from environment variables, defaulting to "pytest".
    pytest_cmd = os.environ.get("LLM_PYTEST_CMD") or os.environ.get("PYTEST_CMD") or "pytest"

    # Define the directory for conversation logs.
    conv_dir = "/home/ubuntu/app/logs"
    # Create the log directory if it doesn't exist.
    os.makedirs(conv_dir, exist_ok=True)
    
    # Generate an ISO-formatted timestamp for the log file name.
    ts = datetime.utcnow().isoformat(timespec="seconds").replace(":", "-")
    # Define the path for the specific log file using the timestamp.
    conv_file = os.path.join(conv_dir, f"llm-setup-{ts}.log")
    # Define the path for a symbolic link that points to the latest log file.
    symlink_path = os.path.join(conv_dir, "llm-setup.log")
    
    # --- Symlink Creation ---
    try:
        # Check if the symlink path exists (either as a link or a regular file/dir).
        if os.path.islink(symlink_path) or os.path.exists(symlink_path):
            try:
                # Attempt to delete the existing symlink or file/directory.
                os.unlink(symlink_path)
            except Exception:
                # Ignore errors during unlinking.
                pass
        # Create a new symbolic link from symlink_path to conv_file.
        os.symlink(conv_file, symlink_path)
    except Exception:
        # Ignore errors during symlink creation (e.g., permission issues).
        pass
    # --- End Symlink Creation ---

    # Derive repo slug for memory files from GITHUB_REPO_URL if available, else work repo basename
    repo_url_for_slug = os.environ.get("GITHUB_REPO_URL", "").strip()
    slug = None
    if repo_url_for_slug:
        parts = repo_url_for_slug.rstrip("/").split("/")
        if len(parts) >= 2:
            org = parts[-2]
            name = parts[-1].replace(".git", "")
            slug = f"{org}-{name}"
    if not slug:
        slug = os.path.basename(work_repo.rstrip("/")) or "repo"

    # Session JSONL file for full conversation logging
    conv_jsonl = os.path.join(conv_dir, f"llm-setup-{ts}.jsonl")
    # Persistent memory JSONL across runs for this repo
    memory_file = os.path.join(conv_dir, f"llm-memory-{slug}.jsonl")

    def _now_iso() -> str:
        return datetime.utcnow().isoformat(timespec="seconds") + "Z"

    def _append_jsonl(path: str, obj: dict) -> None:
        try:
            with open(path, "a", encoding="utf-8") as jf:
                jf.write(json.dumps(obj, ensure_ascii=False) + "\n")
        except Exception:
            # JSONL logging errors should never crash the helper
            pass

    def _load_memory_tail(path: str, max_messages: int = 10) -> list[dict]:
        try:
            if not os.path.exists(path):
                return []
            with open(path, "r", encoding="utf-8") as jf:
                lines = jf.readlines()[-max_messages:]
            out = []
            for line in lines:
                line = line.strip()
                if not line:
                    continue
                try:
                    obj = json.loads(line)
                    role = obj.get("role")
                    content = obj.get("content")
                    if isinstance(role, str) and isinstance(content, str):
                        out.append({"role": role, "content": content})
                except Exception:
                    continue
            return out
        except Exception:
            return []

    def log(msg: str) -> None:
        # Inner function to append a message to the current log file.
        with open(conv_file, "a", encoding="utf-8") as f:
            # Open the log file in append mode ('a') with UTF-8 encoding.
            f.write(f"[{_now_iso()}] {msg}\n")
            # Write the message with timestamp followed by a newline.

    # Initial logging of the helper start and key parameters.
    log("=== LLM Setup Helper Started ===")
    log(f"Repo: {work_repo}")
    log(f"Tests root: {tests_root}")
    log(f"Pytest cmd: {pytest_cmd}")
    _append_jsonl(conv_jsonl, {"ts": _now_iso(), "event": "start", "repo": work_repo, "tests_root": tests_root, "pytest_cmd": pytest_cmd})

    # Generate the file tree string for the repository.
    tree = list_tree(work_repo)

    # Prepare Anthropic client lazily (import only when needed, after API key check).
    try:
        from anthropic import Anthropic
        # Attempt to import the Anthropic client library.
    except Exception as e:
        # If import fails, log the error and exit with an error code.
        log(f"Anthropic client import failed: {e}")
        return 1

    # Initialize the Anthropic client with the retrieved API key.
    client = Anthropic(api_key=api_key)

    # Define the system prompt, which sets the LLM's role, goal, and output format/rules.
    system_prompt = (
        "You are a reliable setup assistant. Your goal: get tests running to at least 50% passing. "
        "You can only respond with JSON following the schema: {\n"
        "  \"actions\": [ { \"type\": \"shell\", \"cmd\": \"string\" } ... ],\n"
        "  \"comment\": \"brief reasoning\"\n} . "
        "Rules:\n- Only use safe package commands (pip install <pkg>), avoid destructive ops.\n"
        "- DO NOT rewrite pyproject.toml or packaging files destructively. Prefer adding missing deps via pip.\n"
        "- Avoid removing files or changing project layout; prefer the simplest fix.\n"
        "- Assume venv is active.\n- Prefer installing missing Python packages explicitly.\n"
        "- Use one or more actions; keep each action minimal.\n"
        "- You are not allowed to change directories (cd), or use semicolons (;) or pipes (|).\n"
        "- You MAY chain multiple allowed install commands using '&&' only. Each side of '&&' must be an allowed install command.\n"
        "- Use the repo venv interpreter for all Python/pip: .venv/bin/python -m pip install <pkgs> (preferred).\n"
        "- Allowed actions: {type: shell, cmd: '.venv/bin/python -m pip install ...' | 'python -m pip install ...' | 'pip install ...' | 'python devscripts/install_deps.py'} | {type: read, path: 'relative/path'} | {type: finish}.\n"
        "- Do NOT run tests (pytest) or benchmarks; we will run tests after you finish.\n"
        "- After you are done with installation steps, respond with a single {\"actions\":[{\"type\":\"finish\"}],\"comment\":\"...\"}.\n"
        "- We will provide you the remaining tries_left each turn; optimize your plan accordingly.\n"
        "- If you need to read a small file, request it using: {\"actions\":[{\"type\":\"read\",\"path\":\"relative/path\"}]} (limit to key config/tests).\n"
    )

    # Prepare the initial context data to send to the LLM.
    user_context = {
        "repo_path": work_repo,
        "tests_root": tests_root,
        "pytest_cmd": pytest_cmd,
        "tree": tree,
    }

    # Initialize the conversation history with memory, then initial context message.
    chat = []
    prior_memory = _load_memory_tail(memory_file, max_messages=10)
    if prior_memory:
        chat.extend(prior_memory)
        log(f"Loaded conversation memory: {len(prior_memory)} messages")
        _append_jsonl(conv_jsonl, {"ts": _now_iso(), "event": "memory_loaded", "count": len(prior_memory)})
    # Include minimal test tail if available
    test_tail = ""
    try:
        test_log_path = os.environ.get("TEST_LOG_FILE", "").strip()
        if test_log_path and os.path.isfile(test_log_path):
            with open(test_log_path, "r", encoding="utf-8", errors="ignore") as tf:
                content = tf.read()
                test_tail = content[-4000:]
    except Exception:
        test_tail = ""
    initial_payload = {"context": user_context}
    if test_tail:
        initial_payload["last_test_output_tail"] = test_tail
    initial_payload["tries_left"] = int(os.environ.get("LLM_SETUP_MAX_STEPS", "10") or "10")
    chat.append({"role": "user", "content": json.dumps(initial_payload, ensure_ascii=False)})

    # Initialize a set to track executed commands for deduplication (guardrail).
    executed_cmds: set[str] = set()

    def add_msg(role: str, content: str) -> None:
        # Inner function to add a new message (from user or assistant) to the chat history.
        chat.append({"role": role, "content": content})
        # Persist to session JSONL and memory JSONL
        payload = {"ts": _now_iso(), "role": role, "content": content}
        _append_jsonl(conv_jsonl, payload)
        _append_jsonl(memory_file, payload)
        # Also write a readable line to plaintext log
        log(f"[{role.upper()}] {content}")
        # Keep recent window to bound context size for the LLM.
        if len(chat) > 12:
            del chat[: len(chat) - 12]

    # Define loop control variables.
    max_loops = int(os.environ.get("LLM_SETUP_MAX_STEPS", "10") or "10")
    pass_threshold = 0.5 # Target passing ratio (50%).
    
    # Start the main loop for setup attempts.
    for i in range(1, max_loops + 1):
        # Log the start of the current step.
        log(f"\n--- Setup step {i}/{max_loops} ---")
        _append_jsonl(conv_jsonl, {"ts": _now_iso(), "event": "setup_step", "iteration": i, "max": max_loops})
        
        # Prepare the message content for the LLM, including remaining tries and optional signals from last executions.
        msg_text = json.dumps({
            "context": user_context,
            "tries_left": max_loops - i + 1,
        }, ensure_ascii=False)
        
        # Add the new message to the conversation history.
        add_msg("user", msg_text)
        
        # --- Call Anthropic API ---
        try:
            # Get the model name from the environment, with a default value.
            model = os.environ.get("ANTHROPIC_MODEL", "claude-3-5-haiku-20241022")
            # Provide summarized last outputs to reduce tokens
            # (already appended test output tail above; keep conversation short via truncation in add_msg)
            resp = client.messages.create(
                model=model,
                max_tokens=800, # Limit the LLM's response length.
                system=system_prompt, # Provide the system prompt.
                messages=chat, # Send the conversation history.
            )
        except Exception as e:
            # Log API error and exit with an error code.
            log(f"Anthropic API error: {e}")
            _append_jsonl(conv_jsonl, {"ts": _now_iso(), "event": "api_error", "error": str(e)})
            return 1

        # Extract the text content from the LLM response.
        text = ""
        try:
            # Iterate through content blocks (which should contain the JSON response).
            for block in resp.content:
                if getattr(block, "type", "text") == "text":
                    text += block.text
        except Exception:
            # Fallback for unexpected response structure.
            text = str(resp)
            
        # Log the raw text response from the LLM.
        log("LLM raw response:\n" + text)
        add_msg("assistant", text)
        # --- End Anthropic API Call ---

        # Handle optional file read requests from assistant (single small file)
        try:
            data_peek = json.loads(text)
            reqs = data_peek.get("actions") if isinstance(data_peek, dict) else None
            read_path = None
            if isinstance(reqs, list):
                for a in reqs:
                    if isinstance(a, dict) and a.get("type") == "read" and isinstance(a.get("path"), str):
                        read_path = a["path"]
                        break
            if read_path:
                # Allowlist: only small text files under repo root
                safe = True
                if read_path.startswith("/") or ".." in read_path or len(read_path) > 200:
                    safe = False
                abs_path = os.path.join(work_repo, read_path)
                content = ""
                if safe and os.path.isfile(abs_path) and os.path.getsize(abs_path) <= 200*1024:
                    try:
                        with open(abs_path, "r", encoding="utf-8", errors="ignore") as rf:
                            content = rf.read()
                    except Exception:
                        content = ""
                # Return file content (summarized if large) as a user message
                if content:
                    snippet = content if len(content) <= 8000 else (content[:6000] + "\n...\n" + content[-1000:])
                    add_msg("user", json.dumps({"file": read_path, "content": snippet}, ensure_ascii=False))
                else:
                    add_msg("user", json.dumps({"file": read_path, "error": "unavailable or too large"}, ensure_ascii=False))
                # Continue to next loop iteration to let model act on the file
                continue
        except Exception:
            pass

        # --- Parse LLM Response and Execute Actions ---
        # Robust JSON extraction: find the largest JSON object if raw isn't valid
        def _extract_json(s: str) -> str:
            try:
                json.loads(s)
                return s
            except Exception:
                pass
            # best-effort bracket matching
            start = s.find('{')
            last = s.rfind('}')
            while start != -1 and last != -1 and last > start:
                candidate = s[start:last+1]
                try:
                    json.loads(candidate)
                    return candidate
                except Exception:
                    last = s.rfind('}', 0, last)
            return ""

        extracted = _extract_json(text)
        if not extracted:
            log("Failed to extract JSON from assistant output")
            _append_jsonl(conv_jsonl, {"ts": _now_iso(), "event": "parse_error", "error": "no_json_in_output", "assistant_text": text[:16000]})
            return 1

        try:
            data = json.loads(extracted)
        except Exception as e:
            log(f"Failed to parse LLM JSON: {e}")
            _append_jsonl(conv_jsonl, {"ts": _now_iso(), "event": "parse_error", "error": str(e), "assistant_text": text[:16000]})
            return 1

        # Extract the 'actions' list from the parsed JSON data.
        actions = data.get("actions") or []
        # Validate that 'actions' is a non-empty list.
        if not isinstance(actions, list) or not actions:
            log("No actions provided; aborting.")
            _append_jsonl(conv_jsonl, {"ts": _now_iso(), "event": "no_actions"})
            return 1

        # Detect finish action
        if len(actions) == 1 and isinstance(actions[0], dict) and actions[0].get("type") == "finish":
            log("Assistant requested finish; exiting setup loop.")
            _append_jsonl(conv_jsonl, {"ts": _now_iso(), "event": "finish_requested", "iteration": i})
            # Successful finish; orchestrator will re-run tests
            return 0

        # List to store the summary of executed commands for feedback to the LLM.
        exec_summ = []
        # Iterate over the proposed actions.
        for a in actions:
            # Validate the action structure
            if not isinstance(a, dict):
                continue
            a_type = a.get("type")
            if a_type == "shell":
                # Extract and clean the command string.
                cmd = str(a.get("cmd") or "").strip()
                if not cmd:
                    continue
                
                # Guardrails: determine if this shell command is allowed
                def _is_allowed_single(c: str) -> bool:
                    lc = c.strip()
                    # Disallow directory changes
                    if lc.startswith("cd ") or " cd " in f" {lc} ":
                        return False
                    # Allow pip installs via multiple forms
                    allowed_prefixes = [
                        "pip install ",
                        "pip3 install ",
                        "python -m pip install ",
                        "python3 -m pip install ",
                        ".venv/bin/python -m pip install ",
                        ".venv/bin/python3 -m pip install ",
                    ]
                    if any(lc.startswith(p) for p in allowed_prefixes):
                        return True
                    # Allow safe pip queries
                    allowed_pip_queries = [
                        "python -m pip list",
                        "python3 -m pip list",
                        ".venv/bin/python -m pip list",
                        ".venv/bin/python3 -m pip list",
                        "python -m pip show ",
                        "python3 -m pip show ",
                        ".venv/bin/python -m pip show ",
                        ".venv/bin/python3 -m pip show ",
                        "python -m pip freeze",
                        "python3 -m pip freeze",
                        ".venv/bin/python -m pip freeze",
                        ".venv/bin/python3 -m pip freeze",
                    ]
                    if any(lc.startswith(p) for p in allowed_pip_queries):
                        return True
                    # Allow repo-provided dependency helper when present
                    if lc in {
                        "python devscripts/install_deps.py",
                        "python3 devscripts/install_deps.py",
                        ".venv/bin/python devscripts/install_deps.py",
                        ".venv/bin/python3 devscripts/install_deps.py",
                    }:
                        return True
                    return False
                
                def _is_allowed_shell(c: str) -> bool:
                    # Disallow semicolons or pipes entirely
                    if ";" in c or "|" in c:
                        return False
                    # Permit '&&' as a chain of allowed singles
                    parts = [p.strip() for p in c.split("&&")]
                    if len(parts) > 1:
                        return all(_is_allowed_single(p) for p in parts)
                    return _is_allowed_single(c)
                
                if not _is_allowed_shell(cmd):
                    log(f"Blocked command (unsafe): {cmd}")
                    exec_summ.append({"cmd": cmd, "status": "blocked"})
                    continue
                if cmd in executed_cmds:
                    log(f"Skipping duplicate command: {cmd}")
                    exec_summ.append({"cmd": cmd, "status": "duplicate_skipped"})
                    continue
                # Execute the command in repo cwd
                old_cwd = os.getcwd()
                try:
                    os.chdir(work_repo)
                    log(f"Executing: {cmd}")
                    r = sh(cmd)
                    log(r.stdout)
                    _append_jsonl(conv_jsonl, {"ts": _now_iso(), "event": "exec", "cmd": cmd, "stdout": r.stdout[-16000:]})
                    executed_cmds.add(cmd)
                    exec_summ.append({"cmd": cmd, "status": "ran"})
                finally:
                    os.chdir(old_cwd)
            elif a_type == "read":
                # handled earlier in the loop; ignore here
                continue
            elif a_type == "finish":
                log("Assistant issued finish in multi-action response; exiting.")
                _append_jsonl(conv_jsonl, {"ts": _now_iso(), "event": "finish_requested", "iteration": i})
                return 0
            else:
                # Unknown type
                exec_summ.append({"type": a_type or "", "status": "blocked"})
        # Feed back the execution summary to the LLM for memory/context in the next loop.
        if exec_summ:
            add_msg("user", json.dumps({"executed": exec_summ}, ensure_ascii=False))

    # If we exhaust steps without finish, return non-zero so orchestrator can decide next round
    log("Setup steps exhausted without finish action.")
    _append_jsonl(conv_jsonl, {"ts": _now_iso(), "event": "exhausted"})
    return 1


if __name__ == "__main__":
    # Standard entry point of the script. Calls main() and uses its return value as the process exit code.
    sys.exit(main())