mega_coder.py

#!/usr/bin/env python3
import io
import os
import re
import sys
import time
import random
import subprocess
import warnings

from datetime import datetime

from typing import Optional

from dotenv import load_dotenv
from colorama import Fore, Style, init as colorama_init
from tqdm import tqdm

import google.generativeai as genai

try:
    from gitingest import ingest
except ImportError:
    ingest = None  # type: ignore[assignment]

# screen capture dependencies
try:
    import mss
    import numpy as np
except ImportError:
    mss = None
    np = None

try:
    from openai import OpenAI
except ImportError:
    OpenAI = None

load_dotenv()
colorama_init(autoreset=True)
warnings.filterwarnings("ignore", category=DeprecationWarning)

APP_TITLE = "I’m Mega Coder. What would you like me to do today?"
MENU = (
    "1. Develop a python program.\n"
    "2. Fix/change something in a Github repository.\n"
    "3. Look at my screen and give me realtime coding tips."
)

GENERATED_FILE = "generated-code-gemini.py"
GEMINI_MODEL = "gemini-2.5-flash-lite"
GEMINI_PRO_MODEL = "gemini-2.5-pro"
GPT5_MINI_MODEL = "gpt-5-mini"

RUN_LOG_FILE = "mega_coder.log"


# Max characters from the repository digest that we pass to the LLM.
# This is to avoid creating an extremely huge prompt.
MAX_REPO_DIGEST_CHARS = 200_000

SYSTEM_INSTRUCTION = (
    "You are a Python code generator. Always output only complete Python source code, "
    "with no markdown fences or explanations. Never use command-line arguments "
    "(sys.argv, argparse, click, typer) and never call input() or wait for user input. "
    "The program must run end-to-end with zero stops. "
    "Use only Python's standard library; do not import any third-party packages. "
    "Include meaningful assert statements that validate the core logic of the program, "
    "not trivial truths like 'assert 1 == 1'. "
    "If an assert fails, let the AssertionError propagate so the process exits with a non-zero status."
)

REPO_SYSTEM_INSTRUCTION = (
    "You are an expert software engineer. You receive a text digest of a Git repository "
    "and a natural language request from the user. Use only the information in the digest "
    "to reason about the repository. When appropriate, point to concrete file paths and "
    "symbols. If something is impossible to know from the digest, say that explicitly."
)

MAX_FIX_ATTEMPTS = 5
MAX_LINT_FIX_ATTEMPTS = 3

MUTATION_PROBABILITY = 0.3

SCREEN_GRAB_INTERVAL_SECONDS = 1.0

ANSI_ESCAPE_RE = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")
CTRL_CHARS_RE = re.compile(r"[\x00-\x1F\x7F]")

FORBIDDEN_PATTERNS = [
    r"\binput\s*\(",
    r"\bsys\.argv\b",
    r"\bargparse\b",
    r"\bclick\b",
    r"\btyper\b",
]


def ensure_no_cli_args() -> None:
    """
    Enforce "no command-line arguments" requirement.
    Exits with code 2 if arguments are provided.
    """
    if len(sys.argv) > 1:
        print(
            Fore.RED
            + "This app does not accept command-line arguments. Please run without arguments."
            + Style.RESET_ALL
        )
        sys.exit(2)


def safe_input(prompt: str) -> Optional[str]:
    """
    Read a single line safely.
    Returns stripped string, or None if EOF or Ctrl+C is encountered.
    """
    try:
        s = input(prompt)
        return s.strip()
    except (EOFError, KeyboardInterrupt):
        return None


def sanitize_user_description(text: str) -> str:
    """
    Remove ANSI escape sequences and control characters from user input.
    Also collapse excessive whitespace and strip leading and trailing spaces.
    """
    if text is None:
        return ""
    cleaned = ANSI_ESCAPE_RE.sub("", text)
    cleaned = CTRL_CHARS_RE.sub("", cleaned)
    cleaned = re.sub(r"\s+", " ", cleaned).strip()
    return cleaned


def is_meaningful_description(text: str, min_len: int = 10) -> bool:
    """
    Heuristic: require some letters or digits and a minimum length after sanitization.
    """
    if not text or len(text) < min_len:
        return False
    return bool(re.search(r"[A-Za-z0-9]", text))


def read_program_description() -> Optional[str]:
    """
    Repeatedly ask the user for a meaningful program description.
    Returns the cleaned description, or None if the user cancels and wants
    to go back to the main menu.
    Exits the app on EOF or Ctrl+C.
    """
    while True:
        desc = safe_input("\nYour description (or 'q' to go back to the main menu): ")
        if desc is None:
            print("No input (EOF or Ctrl+C). Exiting.")
            sys.exit(0)

        if desc.lower() in {"q", "quit", "exit"}:
            print("Cancelled by user, returning to the main menu.")
            return None

        desc = sanitize_user_description(desc)

        if not desc:
            print(
                "Description is empty. Please describe the Python program you want me to develop, "
                "or type 'q' to cancel."
            )
            continue

        if not is_meaningful_description(desc):
            print(
                "Description is not meaningful. Please give a clearer description of the Python program you want, "
                "or type 'q' to cancel."
            )
            continue

        return desc


def read_github_repo_url() -> Optional[str]:
    """
    Ask the user for a public GitHub repository URL for gitingest.
    Returns the URL or None if the user cancels and wants to go back to the menu.
    """
    while True:
        url = safe_input(
            "\nGive me the full url of a public github repository (or 'q' to go back to the main menu): "
        )
        if url is None:
            print("No input (EOF or Ctrl+C). Exiting.")
            sys.exit(0)

        url = url.strip()
        if url.lower() in {"q", "quit", "exit"}:
            print("Cancelled by user, returning to the main menu.")
            return None

        if not url:
            print("Repository url is empty. Please paste a full GitHub url.")
            continue

        if "github.com" not in url:
            print("This does not look like a GitHub url. Please paste a url that contains 'github.com'.")
            continue

        return url


def read_repo_request() -> Optional[str]:
    """
    Ask the user what to fix, change, or explain in the selected repository.
    Returns the cleaned text or None if the user cancels.
    """
    while True:
        text = safe_input(
            "\nTell me what you want me to fix/change/explain in that repository (or 'q' to go back to the main menu): "
        )
        if text is None:
            print("No input (EOF or Ctrl+C). Exiting.")
            sys.exit(0)

        if text.lower() in {"q", "quit", "exit"}:
            print("Cancelled by user, returning to the main menu.")
            return None

        text = sanitize_user_description(text)
        if not text:
            print(
                "Please describe what you want me to fix, change, or explain in that repository, "
                "or type 'q' to cancel."
            )
            continue

        return text


def read_menu_choice() -> Optional[int]:
    """
    Prompt until a valid menu choice 1..3 is given, or return None to exit.
    Accepts 'q' or 'quit' or 'exit' for a graceful exit.
    """
    while True:
        line = safe_input("\nEnter your choice (1, 2, or 3) or 'q' to quit: ")
        if line is None:
            return None
        if line.lower() in {"q", "quit", "exit"}:
            return None
        if not line:
            print("Please enter a choice: 1, 2, or 3 (or 'q' to quit).")
            continue
        if not line.isdigit():
            print("Invalid input. Please enter a number: 1, 2, or 3.")
            continue
        choice = int(line)
        if choice not in (1, 2, 3):
            print("Out of range. Please choose 1, 2, or 3.")
            continue
        return choice


def build_gemini_prompt(user_description: str) -> str:
    """
    Build a strict prompt for Gemini: no CLI args, no input, runnable end to end,
    code only, with asserts.
    """
    return f"""
You are an expert Python code generator. Create one complete, runnable Python 3 program.

Hard requirements:
- Do NOT use command-line arguments (assume none will ever be passed).
- Do NOT call input() or wait for user input. The program must run end-to-end with zero stops.
- Output ONLY raw Python source code. Do NOT include explanations or markdown fences.
- Include a short module-level docstring at the top describing the program.
- Use only Python's standard library; do not import any external or third-party packages.
- Include meaningful assert statements that validate the key logic of the program and avoid trivial checks
  (do not write asserts like 'assert 1 == 1'). If something in the logic is wrong, these asserts must raise
  an AssertionError.
- Follow PEP8 where reasonable.

User request:
\"\"\"{user_description}\"\"\"

Return ONLY the final Python source code. No backticks. No prose. No markdown.
"""


def _generate_with_gemini_model(
    model_name: str,
    prompt: str,
    system_instruction: Optional[str] = None,
) -> str:
    """
    Low level helper that calls a Gemini model and returns the raw text response.
    Used by both the Flash Lite and Pro flows.
    """
    api_key = os.getenv("GOOGLE_API_KEY")
    if not api_key:
        print("Missing GOOGLE_API_KEY environment variable. Please export it and try again.")
        sys.exit(1)

    try:
        genai.configure(api_key=api_key)
        if system_instruction:
            model = genai.GenerativeModel(
                model_name,
                system_instruction=system_instruction,
            )
        else:
            model = genai.GenerativeModel(model_name)
        response = model.generate_content(prompt)
    except Exception as exc:
        print(f"Gemini API error: {exc}")
        sys.exit(1)

    text = getattr(response, "text", None)
    if not text:
        try:
            if response.candidates and response.candidates[0].content.parts:
                parts = []
                for part in response.candidates[0].content.parts:
                    t = getattr(part, "text", None)
                    if t:
                        parts.append(t)
                text = "".join(parts)
        except Exception:
            text = None

    if not text or not text.strip():
        print("Gemini returned an empty response. Please try a different description or refine the prompt.")
        sys.exit(1)

    return text.strip()


def generate_with_gemini(prompt: str) -> str:
    """
    Wrapper for Gemini 2.5 Flash Lite that you already use for option 1.
    """
    return _generate_with_gemini_model(GEMINI_MODEL, prompt, SYSTEM_INSTRUCTION)


def generate_with_gemini_pro(prompt: str) -> str:
    """
    Wrapper for Gemini 2.5 Pro used for repository analysis in option 2.
    """
    return _generate_with_gemini_model(GEMINI_PRO_MODEL, prompt, REPO_SYSTEM_INSTRUCTION)


def extract_python_code(raw_text: str) -> str:
    """
    Extract Python code from the model response.
    If fenced blocks exist, prefer the first python-labeled block, else return trimmed text.
    """
    fence = re.search(r"```python\s*(.*?)```", raw_text, flags=re.DOTALL | re.IGNORECASE)
    if fence:
        return fence.group(1).strip()
    fence_any = re.search(r"```\s*(.*?)```", raw_text, flags=re.DOTALL)
    if fence_any:
        return fence_any.group(1).strip()
    return raw_text.strip()


def build_repo_digest(repo_url: str) -> str:
    """
    Use gitingest to turn a GitHub repository into a single text digest for LLM consumption.
    """
    if ingest is None:
        raise RuntimeError(
            "gitingest is not installed. Please run 'pip install gitingest' in your virtualenv and try again."
        )

    try:
        print(
            Fore.CYAN
            + "\n[gitingest] Downloading and analyzing repository, this can take a while..."
            + Style.RESET_ALL
        )
        summary, tree, content = ingest(
            repo_url,
            include_submodules=False,
        )
    except Exception as exc:
        raise RuntimeError(f"gitingest failed for url '{repo_url}': {exc}") from exc

    parts: list[str] = []
    if summary:
        parts.append("### SUMMARY ###\n")
        parts.append(str(summary).strip())
        parts.append("\n\n")
    if tree:
        parts.append("### TREE ###\n")
        parts.append(str(tree).strip())
        parts.append("\n\n")
    if content:
        parts.append("### CONTENT ###\n")
        parts.append(str(content).strip())

    digest = "".join(parts).strip()

    if not digest:
        raise RuntimeError("gitingest returned an empty digest for this repository")

    if len(digest) > MAX_REPO_DIGEST_CHARS:
        print(
            Fore.YELLOW
            + f"\n[gitingest] Digest is very large, keeping only the first {MAX_REPO_DIGEST_CHARS} characters for the LLM."
            + Style.RESET_ALL
        )
        digest = (
            digest[:MAX_REPO_DIGEST_CHARS]
            + "\n\n[Digest truncated in mega_coder.py, repository is larger than this snippet.]"
        )

    return digest


def build_repo_prompt(repo_url: str, repo_digest: str, user_request: str) -> str:
    """
    Build the prompt that will be sent to Gemini 2.5 Pro for repository analysis.
    """
    return f"""
You are an expert software engineer and code reviewer.

You are given:
1. A GitHub repository url.
2. A text digest of the repository that was produced by gitingest.
3. A concrete request from the user about what to fix, change, or explain.

Important rules:
- Base your answer only on the information visible in the digest.
- If the digest does not contain enough information to answer something, say so clearly.
- When you refer to code, mention file names and function or class names when possible.
- If the user asked for a fix, describe the change in detail and provide example patches.
- If the user asked for an explanation, explain step by step and keep it practical.

GITHUB REPOSITORY URL:
{repo_url}

REPOSITORY DIGEST (gitingest):
<<<REPO_DIGEST_START>>>
{repo_digest}
<<<REPO_DIGEST_END>>>

USER REQUEST:
<<<USER_REQUEST_START>>>
{user_request}
<<<USER_REQUEST_END>>>
"""


def write_generated_file(code: str, path: str) -> None:
    """
    Write generated code to a UTF-8 file. Ensure trailing newline for POSIX friendliness.
    """
    with open(path, "w", encoding="utf-8") as f:
        f.write(code if code.endswith("\n") else code + "\n")


def violates_io_constraints(code: str) -> bool:
    """
    Return True if the code violates the no-input and no-CLI-args constraints.
    """
    return any(re.search(pat, code) for pat in FORBIDDEN_PATTERNS)


def run_generated_code(path: str, timeout: int = 20) -> tuple[int, str, str]:
    """
    Run the generated Python program as a separate process using the same interpreter.
    Returns a tuple (exit_code, stdout, stderr).
    """
    script_path = os.path.abspath(path)
    print(Fore.BLUE + f"\n[runtime] Running generated program: {script_path}" + Style.RESET_ALL)

    cmd = [sys.executable, script_path]

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=timeout,
            check=False,
        )
        return result.returncode, result.stdout, result.stderr

    except subprocess.TimeoutExpired as exc:
        stdout = exc.stdout or ""
        stderr = (exc.stderr or "") + f"\nTimeout after {timeout} seconds."
        return 124, stdout, stderr

    except FileNotFoundError:
        return 127, "", f"File not found: {script_path}"

    except PermissionError:
        return 126, "", f"Permission denied: {script_path}"

    except OSError as exc:
        return 1, "", f"OS error while running '{script_path}': {exc}"


def maybe_mutate_code(code: str) -> tuple[str, bool]:
    """
    With probability MUTATION_PROBABILITY, inject a tiny, deterministic syntax bug.
    Returns (code_to_write, mutated_flag).
    """
    if random.random() < MUTATION_PROBABILITY:
        buggy = code.rstrip() + "\nprint('INTENTIONAL BUG'\n"
        print(
            Fore.MAGENTA
            + f"\n[mutator] Injected a tiny syntax bug into '{GENERATED_FILE}' on purpose, "
            "just to keep the fixer busy."
            + Style.RESET_ALL
        )
        return buggy, True

    print(
        Fore.MAGENTA
        + f"\n[mutator] Wrote '{GENERATED_FILE}' without mutation this time."
        + Style.RESET_ALL
    )
    return code, False


def print_run_summary(exit_code: int, stdout: str, stderr: str) -> None:
    """
    Pretty-print a summary of a single run of the generated program.
    """
    print("\n=== Run summary ===")
    print(f"Exit code: {exit_code}")

    if stdout:
        print("\nProgram output:")
        print("---------------")
        print(stdout.rstrip())

    if stderr:
        print("\nProgram errors:")
        print("---------------")
        print(stderr.rstrip())

    if exit_code == 0:
        print(
            "\nResult: "
            + Fore.GREEN
            + "success"
            + Style.RESET_ALL
            + ", the program finished with exit code 0."
        )
        print(
            "Assuming the asserts are meaningful, this suggests that the logic is consistent."
        )
    else:
        print(
            "\nResult: "
            + Fore.RED
            + "failure"
            + Style.RESET_ALL
            + ", the program did not finish cleanly. I will try to auto-fix it if possible."
        )


def build_fix_prompt(
    user_description: str,
    buggy_code: str,
    exit_code: int,
    stdout: str,
    stderr: str,
) -> str:
    """
    Build a prompt that asks Gemini to fix the generated code, given
    the error details from the last run.
    """
    return f"""
You generated a Python 3 program for the following user request:

USER REQUEST:
\"\"\"{user_description}\"\"\"

When I ran the program, it failed with the following details:

EXIT CODE:
{exit_code}

STDOUT:
<<<STDOUT START>>>
{stdout}
<<<STDOUT END>>>

STDERR:
<<<STDERR START>>>
{stderr}
<<<STDERR END>>>

Here is the exact code that produced the error:

CURRENT CODE:
<<<CODE START>>>
{buggy_code}
<<<CODE END>>>

Please return a corrected FULL version of the same program that:

- Keeps the same functionality and overall structure.
- Preserves or improves the assert statements so they still validate the core logic.
- Does NOT use command-line arguments (no sys.argv, argparse, click, typer).
- Does NOT call input() or wait for user input.
- Uses only the Python standard library.
- Runs end-to-end without stopping for user input.

Return ONLY the complete corrected Python source code. No backticks. No prose. No markdown.
"""


def fix_code_with_gemini(
    current_code: str,
    exit_code: int,
    stdout: str,
    stderr: str,
    user_description: str,
) -> str:
    """
    Ask Gemini to fix the generated program based on the last run result.
    Returns the fixed code, or the original code if the response is invalid.
    """
    prompt = build_fix_prompt(user_description, current_code, exit_code, stdout, stderr)
    print(Fore.CYAN + "\n[auto-fix] Asking Gemini to fix the program based on the error report..." + Style.RESET_ALL)
    raw_text = generate_with_gemini(prompt)
    fixed_code = extract_python_code(raw_text)

    if not fixed_code or (
        "def " not in fixed_code
        and "class " not in fixed_code
        and "import " not in fixed_code
    ):
        print(
            Fore.YELLOW
            + "[auto-fix] Gemini response does not look like valid Python code. "
            "Keeping the previous version."
            + Style.RESET_ALL
        )
        return current_code

    if violates_io_constraints(fixed_code):
        print(
            Fore.YELLOW
            + "[auto-fix] Gemini produced code that violates the IO constraints. "
            "Keeping the previous version."
            + Style.RESET_ALL
        )
        return current_code

    print(Fore.GREEN + "[auto-fix] Gemini returned a candidate fixed version of the program." + Style.RESET_ALL)
    return fixed_code


def auto_fix_until_runs(initial_code: str, user_description: str) -> Optional[str]:
    """
    Try to run and auto-fix the code up to MAX_FIX_ATTEMPTS times.
    Returns the final working version of the code, or None if we give up.
    """
    current_code = initial_code

    print(
        Fore.CYAN
        + "\n=== Step 1/3: generate and auto-fix the program ==="
        + Style.RESET_ALL
    )

    for attempt in range(1, MAX_FIX_ATTEMPTS + 1):
        print(
            Fore.CYAN
            + f"\n---------- Auto fix attempt {attempt}/{MAX_FIX_ATTEMPTS} ----------"
            + Style.RESET_ALL
        )

        code_to_write, mutated = maybe_mutate_code(current_code)
        try:
            write_generated_file(code_to_write, GENERATED_FILE)
        except OSError as exc:
            print(Fore.RED + f"Failed to write '{GENERATED_FILE}': {exc}" + Style.RESET_ALL)
            return None

        if mutated:
            print(Fore.BLUE + "[runner] This run uses a deliberately mutated version of the program." + Style.RESET_ALL)
        else:
            print(Fore.BLUE + "[runner] This run uses the clean version of the program." + Style.RESET_ALL)

        exit_code, stdout, stderr = run_generated_code(GENERATED_FILE)
        print_run_summary(exit_code, stdout, stderr)

        if exit_code == 0:
            return code_to_write

        if attempt == MAX_FIX_ATTEMPTS:
            return None

        current_code = fix_code_with_gemini(
            code_to_write, exit_code, stdout, stderr, user_description
        )


def run_pylint(path: str, timeout: int = 30) -> tuple[int, str]:
    """
    Run pylint on the given file.
    Returns (exit_code, full_output). Exit code 0 means no issues.
    We disable C and R so we focus on warnings and errors, not style hints.
    """
    cmd = ["pylint", "--disable=C,R", path]

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=timeout,
            check=False,
        )
    except subprocess.TimeoutExpired as exc:
        output = (exc.stdout or "") + "\n[lint] pylint timed out."
        return 124, output
    except FileNotFoundError:
        return 127, "[lint] pylint command not found. Is it installed?"
    except OSError as exc:
        return 1, f"[lint] OS error while running pylint: {exc}"

    output = (result.stdout or "") + (result.stderr or "")
    return result.returncode, output


def build_lint_fix_prompt(
    user_description: str,
    current_code: str,
    pylint_output: str,
) -> str:
    """
    Build a prompt that asks Gemini to fix lint warnings and errors from pylint.
    """
    return f"""
You previously wrote this Python 3 program for the following user request:

USER REQUEST:
\"\"\"{user_description}\"\"\"

CURRENT CODE:
<<<CODE START>>>
{current_code}
<<<CODE END>>>

I ran pylint (with C and R messages disabled) on this file and got:

PYLINT OUTPUT:
<<<PYLINT START>>>
{pylint_output}
<<<PYLINT END>>>

Please return a new version of the program that:

- Fixes the pylint warnings and errors shown above.
- Keeps the same functionality and overall behavior.
- Preserves the existing assert statements or replaces them with equivalent ones
  that validate the same conditions.
- Does NOT use command-line arguments (no sys.argv, argparse, click, typer).
- Does NOT call input() or wait for user input.
- Uses only the Python standard library.
- Runs end-to-end without stopping for user input.

Return ONLY the full corrected Python source code. No backticks. No prose. No markdown.
"""


def fix_lint_with_gemini(
    current_code: str,
    pylint_output: str,
    user_description: str,
) -> str:
    """
    Ask Gemini to fix pylint warnings and errors, given the current code and
    the pylint output. Returns the fixed code or the original code if invalid.
    """
    prompt = build_lint_fix_prompt(user_description, current_code, pylint_output)
    print(Fore.CYAN + "\n[lint] Asking Gemini to fix pylint issues..." + Style.RESET_ALL)
    raw_text = generate_with_gemini(prompt)
    fixed_code = extract_python_code(raw_text)

    if not fixed_code or (
        "def " not in fixed_code
        and "class " not in fixed_code
        and "import " not in fixed_code
    ):
        print(
            Fore.YELLOW
            + "[lint] Gemini response does not look like valid Python. Keeping the previous version."
            + Style.RESET_ALL
        )
        return current_code

    if violates_io_constraints(fixed_code):
        print(
            Fore.YELLOW
            + "[lint] Gemini produced code that violates the IO constraints. "
            "Keeping the previous version."
            + Style.RESET_ALL
        )
        return current_code

    print(Fore.GREEN + "[lint] Gemini returned a candidate lint-fixed version." + Style.RESET_ALL)
    return fixed_code


def lint_and_auto_fix(initial_code: str, user_description: str) -> str:
    """
    Run pylint up to MAX_LINT_FIX_ATTEMPTS times and try to auto fix issues with Gemini.
    Returns the latest version of the code, possibly updated.
    """
    current_code = initial_code

    try:
        write_generated_file(current_code, GENERATED_FILE)
    except OSError as exc:
        print(
            Fore.RED
            + f"[lint] Failed to write initial code for linting: {exc}"
            + Style.RESET_ALL
        )
        return current_code

    print(
        Fore.CYAN
        + "\n[lint] Starting pylint checks and auto-fix loop..."
        + Style.RESET_ALL
    )

    with tqdm(
        total=MAX_LINT_FIX_ATTEMPTS,
        desc="Lint fixing",
        unit="round",
        leave=False,
    ) as pbar:
        for round_index in range(1, MAX_LINT_FIX_ATTEMPTS + 1):
            print(
                Fore.CYAN
                + f"\n[lint] Lint round {round_index}/{MAX_LINT_FIX_ATTEMPTS} on '{GENERATED_FILE}'..."
                + Style.RESET_ALL
            )

            exit_code, lint_output = run_pylint(GENERATED_FILE)

            pbar.update(1)

            if exit_code == 0:
                print(
                    Fore.GREEN
                    + "[lint] pylint found no warnings/errors."
                    + Style.RESET_ALL
                )
                print("Amazing. No lint errors/warnings")

                if round_index == MAX_LINT_FIX_ATTEMPTS:
                    return current_code
                continue

            print(
                Fore.YELLOW
                + "[lint] pylint reported warnings/errors:"
                + Style.RESET_ALL
            )
            print(lint_output.strip() or "[lint] (no detailed output)")

            if round_index == MAX_LINT_FIX_ATTEMPTS:
                print("There are still lint errors/warnings")
                return current_code

            try:
                with open(GENERATED_FILE, "r", encoding="utf-8") as f:
                    current_code = f.read()
            except OSError as exc:
                print(
                    Fore.RED
                    + f"[lint] Failed to read '{GENERATED_FILE}' before lint-fix: {exc}"
                    + Style.RESET_ALL
                )
                print("There are still lint errors/warnings")
                return current_code

            fixed_code = fix_lint_with_gemini(current_code, lint_output, user_description)
            current_code = fixed_code

            try:
                write_generated_file(current_code, GENERATED_FILE)
            except OSError as exc:
                print(
                    Fore.RED
                    + f"[lint] Failed to write lint-fixed version: {exc}"
                    + Style.RESET_ALL
                )
                print("There are still lint errors/warnings")
                return current_code

    return current_code


def time_program(path: str, timeout: int = 20) -> tuple[int, float, str, str]:
    """
    Run the generated python file and measure how long it takes.
    Returns (exit_code, elapsed_ms, stdout, stderr).
    """
    start = time.perf_counter()
    exit_code, stdout, stderr = run_generated_code(path, timeout=timeout)
    elapsed_ms = (time.perf_counter() - start) * 1000.0
    return exit_code, elapsed_ms, stdout, stderr


def build_optimization_prompt(user_description: str, current_code: str) -> str:
    """
    Build a prompt that asks Gemini to keep the same asserts but create
    a more efficient version of the program.
    """
    return f"""
You previously wrote a Python program for this description:

USER REQUEST:
\"\"\"{user_description}\"\"\"

The program below works and already contains asserts that verify its logic.

Your task:
- Rewrite this program so that it is more efficient and runs faster.
- Keep all existing asserts and the same logical behavior.

Constraints:
- Do NOT use command-line arguments.
- Do NOT call input() or wait for any user input.
- The program must run end-to-end without stopping.
- Use only Python's standard library.
- Respond with Python code only, no explanations or comments.
- Do NOT use Markdown fences.

Here is the current program:
<<<CODE START>>>
{current_code}
<<<CODE END>>>
"""


def optimize_program_runtime(final_working_code: str, user_description: str) -> None:
    """
    Measure the runtime of the current program, then ask Gemini to generate
    a more efficient version with the same asserts. Compare runtimes and
    print the required message if performance improved.
    """
    print(
        Fore.CYAN
        + "\n=== Step 3/3: measure and optimize runtime ==="
        + Style.RESET_ALL
    )

    print("\nMeasuring how long the current program takes to run...")
    before_exit, before_ms, _, _ = time_program(GENERATED_FILE)

    if before_exit != 0:
        print(
            Fore.YELLOW
            + f"\nWarning: the working program failed during the timing step (exit code {before_exit}). "
            "Skipping optimization."
            + Style.RESET_ALL
        )
        return

    before_ms_int = int(before_ms)
    print(Fore.BLUE + f"Current program execution time: {before_ms_int} milliseconds." + Style.RESET_ALL)

    print(Fore.CYAN + "\nAsking Gemini to generate a more efficient version of the program..." + Style.RESET_ALL)
    optimization_prompt = build_optimization_prompt(user_description, final_working_code)
    optimized_raw = generate_with_gemini(optimization_prompt)
    optimized_code = extract_python_code(optimized_raw)

    if not optimized_code or (
        "def " not in optimized_code
        and "class " not in optimized_code
        and "import " not in optimized_code
    ):
        print(
            Fore.YELLOW
            + "\nThe optimized response does not look like valid Python code. Keeping the original program."
            + Style.RESET_ALL
        )
        return
    

    

    if violates_io_constraints(optimized_code):
        print(
            Fore.YELLOW
            + "\nThe optimized code violates the IO constraints (uses input() or CLI args). "
            "Keeping the original program."
            + Style.RESET_ALL
        )
        return

    try:
        write_generated_file(optimized_code, GENERATED_FILE)
    except OSError as exc:
        print(
            Fore.YELLOW
            + f"\nFailed to write optimized version to '{GENERATED_FILE}': {exc} "
            "Keeping the original program."
            + Style.RESET_ALL
        )
        try:
            write_generated_file(final_working_code, GENERATED_FILE)
        except OSError:
            pass
        return

    print("\nMeasuring how long the optimized program takes to run...")
    after_exit, after_ms, _, _ = time_program(GENERATED_FILE)

    if after_exit != 0:
        print(
            Fore.YELLOW
            + f"\nOptimized program did not run successfully (exit code {after_exit}). "
            "Restoring the original program."
            + Style.RESET_ALL
        )
        try:
            write_generated_file(final_working_code, GENERATED_FILE)
        except OSError:
            pass
        return

    after_ms_int = int(after_ms)

    if after_ms < before_ms:
        print(
            f"\nCode running time optimized! It now runs in {after_ms_int} milliseconds, "
            f"while before it was {before_ms_int} milliseconds"
        )
    else:
        print(
            Fore.YELLOW
            + "\nOptimized program did not run faster than the original. Restoring the original program."
            + Style.RESET_ALL
        )
        try:
            write_generated_file(final_working_code, GENERATED_FILE)
        except OSError:
            pass

def print_option1_summary(user_description: str) -> None:
    """
    Print a short summary of what Mega Coder did in option 1.
    This is a user facing summary only.
    """
    print(Fore.CYAN + "\n=== Mega Coder summary ===" + Style.RESET_ALL)
    print(f"User request: {user_description}")
    print("Auto fix: success")
    print("Lint step: completed, see lint output above for details")
    print("Runtime optimization: completed, see timing output above for details")
    print(f"Generated file: {GENERATED_FILE}")
    print("=====================================")

def append_option1_log(user_description: str) -> None:
    """
    Append a short log entry for option 1 to RUN_LOG_FILE.
    This is a very simple text log for debugging and history.
    """
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    try:
        with open(RUN_LOG_FILE, "a", encoding="utf-8") as log_file:
            log_file.write(f"[{timestamp}] Option 1\n")
            log_file.write(f"Description: {user_description}\n")
            log_file.write("Auto fix: success\n")
            log_file.write("Lint step: completed\n")
            log_file.write("Runtime optimization: completed\n")
            log_file.write(f"Generated file: {GENERATED_FILE}\n")
            log_file.write("---\n")
    except OSError:
        print(
            Fore.YELLOW
            + "[log] Failed to write to run log file."
            + Style.RESET_ALL
        )




def normalize_ocr_text(text: str) -> str:
    """
    Normalize OCR text by stripping empty lines and trailing spaces.
    """
    lines = [line.rstrip() for line in text.splitlines() if line.strip()]
    return "\n".join(lines)


def looks_like_code(text: str) -> bool:
    """
    Heuristic to decide whether OCR text looks like source code.
    It checks for common code keywords and symbols.
    Made a bit more permissive so it will still detect code even if there is some noise
    from the terminal or debug logs.
    """
    if not text or not text.strip():
        return False

    # אם יש סימנים מאוד מובהקים של קוד, נחזיר True ישר
    strong_markers = (
        "def ",
        "class ",
        "if __name__ == \"__main__\":",
        "if __name__ == '__main__':",
        "import ",
        "from ",
    )
    if any(marker in text for marker in strong_markers):
        return True

    lines = [line for line in text.splitlines() if line.strip()]
    if len(lines) < 2:
        return False

    indicators = (
        "def ", "class ", "import ", "from ",
        "for ", "while ", "if ", "elif ", "else:",
        "try:", "except", "return ", "print(",
        "{", "}", "function ", "public ", "private ", "void ", "const ",
        "let ", "var ", "#include", "int ", "string ", "float ", "double ",
        "switch ", "case ", "break;", "continue;", "console.log(",
        "println(", "interface ", "implements ", "extends ",
        "struct ", "enum ", "package ", "boolean ", "null", "nil", "None", "bool",
    )

    score = 0.0
    for line in lines:
        stripped = line.strip()
        if any(tok in line for tok in indicators):
            score += 1.0
        if stripped.startswith("#") or stripped.startswith("//"):
            score += 0.5
        if stripped.endswith(":") or stripped.endswith("{") or stripped.endswith("}") or stripped.endswith(";"):
            score += 0.5

    # נוריד קצת את הסף כדי שיחזיר True גם כשיש רעש ברקע
    threshold = max(1.5, len(lines) * 0.2)
    return score >= threshold




def get_openai_client() -> Optional["OpenAI"]:
    """
    Create an OpenAI client for GPT 5 mini.
    Returns None if the package or api key is missing.
    """
    if OpenAI is None:
        print(
            "[screen-helper] The 'openai' package is not installed. "
            "Please install it with 'pip install openai' to use option 3."
        )
        return None

    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        print(
            "[screen-helper] Missing OPENAI_API_KEY environment variable. "
            "Please set it in your .env file to use option 3."
        )
        return None

    try:
        return OpenAI(api_key=api_key)
    except Exception as exc:
        print(f"[screen-helper] Failed to create OpenAI client: {exc}")
        return None

def ask_gpt5_mini_for_code_help(code_text: str) -> str:
    """
    Send the code text from the screen to GPT 5 mini and return its suggestions.
    The response is a plain text message with tips for improvement.
    """
    client = get_openai_client()
    if client is None:
        return ""

    try:
        response = client.chat.completions.create(
            model=GPT5_MINI_MODEL,
            messages=[
                {
                    "role": "system",
                    "content": (
                        "You are a senior software engineer. "
                        "You receive code copied from the user's screen. "
                        "Give short, actionable tips to improve this code. "
                        "Focus on readability, correctness and performance. "
                        "Respond in plain text, without markdown or bullet points."
                    ),
                },
                {
                    "role": "user",
                    "content": (
                        "Here is the code that is currently visible on my screen:\n\n"
                        f"{code_text}\n\n"
                        "Please give me a few concise suggestions to improve it."
                    ),
                },
            ],
        )
    except Exception as exc:
        print(f"[screen-helper] GPT 5 mini request failed: {exc}")
        return ""

    choice = response.choices[0].message
    content = getattr(choice, "content", None)
    return content or ""

def start_screen_coaching_loop() -> None:
    """
    Grab the screen every SCREEN_GRAB_INTERVAL_SECONDS, run OCR with RapidOCR,
    detect whether the text looks like code, and if it does and changed since
    the previous frame, ask GPT 5 mini for tips and print them.
    This loop runs until the user stops the program with Ctrl+C.
    """
    if mss is None or np is None:
        print(
            Fore.RED
            + "[screen-helper] Option 3 requires 'mss' and 'numpy'. "
            "Please install them in your virtualenv."
            + Style.RESET_ALL
        )
        return

    try:
        from rapidocr_openvino import RapidOCR
    except ImportError:
        print(
            Fore.RED
            + "[screen-helper] Option 3 requires the 'rapidocr-openvino' package. "
            "Please install it in your virtualenv."
            + Style.RESET_ALL
        )
        return
    except Exception as exc:
        print(
            Fore.RED
            + f"[screen-helper] Failed to import RapidOCR OpenVINO: {exc}"
            + Style.RESET_ALL
        )
        return

    try:
        ocr_engine = RapidOCR()
    except Exception as exc:
        print(
            Fore.RED
            + f"[screen-helper] Failed to initialize RapidOCR OpenVINO: {exc}"
            + Style.RESET_ALL
        )
        return

    print(
        Fore.CYAN
        + "[screen-helper] Starting screen capture every "
        + f"{SCREEN_GRAB_INTERVAL_SECONDS} seconds. Press Ctrl+C to stop."
        + Style.RESET_ALL
    )

    last_code_text: Optional[str] = None

    try:
        with mss.mss() as sct:
            monitor = sct.monitors[0]

            while True:
                shot = sct.grab(monitor)
                img = np.array(shot)[:, :, :3]

                try:
                    ocr_output = ocr_engine(img)
                except Exception as exc:
                    print(
                        Fore.RED
                        + f"[screen-helper] RapidOCR failed on current frame: {exc}"
                        + Style.RESET_ALL
                    )
                    time.sleep(SCREEN_GRAB_INTERVAL_SECONDS)
                    continue

                if isinstance(ocr_output, tuple):
                    boxes_text_scores = ocr_output[0]
                else:
                    boxes_text_scores = ocr_output

                if not boxes_text_scores:
                    time.sleep(SCREEN_GRAB_INTERVAL_SECONDS)
                    continue

                texts: list[str] = []
                for item in boxes_text_scores:
                    if isinstance(item, (list, tuple)) and len(item) >= 2:
                        texts.append(str(item[1]))

                raw_text = "\n".join(texts)
                normalized = normalize_ocr_text(raw_text)

                if not normalized:
                    time.sleep(SCREEN_GRAB_INTERVAL_SECONDS)
                    continue

                if not looks_like_code(normalized):
                    time.sleep(SCREEN_GRAB_INTERVAL_SECONDS)
                    continue

                if last_code_text is not None and last_code_text == normalized:
                    time.sleep(SCREEN_GRAB_INTERVAL_SECONDS)
                    continue

                last_code_text = normalized

                print(
                    Fore.CYAN
                    + "\n[screen-helper] Detected new code on screen, asking GPT 5 mini for tips..."
                    + Style.RESET_ALL
                )

                tips = ask_gpt5_mini_for_code_help(normalized)
                if tips:
                    print(
                        Fore.MAGENTA
                        + "[screen-helper] Suggestions for the code on screen:"
                        + Style.RESET_ALL
                    )
                    print(tips.strip())
                else:
                    print(
                        Fore.YELLOW
                        + "[screen-helper] GPT 5 mini did not return any suggestions."
                        + Style.RESET_ALL
                    )

                time.sleep(SCREEN_GRAB_INTERVAL_SECONDS)
    except KeyboardInterrupt:
        print(
            Fore.CYAN
            + "\n[screen-helper] Stopped screen capture, returning to main menu."
            + Style.RESET_ALL
        )



def handle_option_1() -> None:
    """
    Option 1:
    ask for a program description,
    generate code with Gemini,
    auto fix it, lint, and optionally optimize runtime.
    """
    print(Fore.CYAN + "\n=== Option 1: develop a Python program ===" + Style.RESET_ALL)
    print("\nDescribe me which python program you want me to develop.")

    desc = read_program_description()
    if desc is None:
        return

    print(Fore.CYAN + "\nGenerating code with Gemini..." + Style.RESET_ALL)
    raw_text = generate_with_gemini(build_gemini_prompt(desc))
    code = extract_python_code(raw_text)

    if not code or ("def " not in code and "class " not in code and "import " not in code):
        print("The response does not look like Python code. Please refine your description and try again.")
        sys.exit(1)

    if violates_io_constraints(code):
        print(
            "Generated code violates the constraints, uses input() or CLI args. "
            "Please refine your description and try again."
        )
        sys.exit(1)

    print(
        Fore.CYAN
        + "\nStarting auto fix loop, I will run the program and ask Gemini to fix it "
        "if it crashes or fails its asserts."
        + Style.RESET_ALL
    )

    final_code = auto_fix_until_runs(code, desc)

    if final_code is None:
        print("\nSorry master, I have failed you. I can’t create this program without issues")
        sys.exit(1)

    try:
        write_generated_file(final_code, GENERATED_FILE)
    except OSError as exc:
        print(Fore.RED + f"Failed to write final version to '{GENERATED_FILE}': {exc}" + Style.RESET_ALL)
        sys.exit(1)

    print(
        Fore.GREEN
        + f"\nAuto fix process completed successfully. The final working version is saved in '{GENERATED_FILE}'."
        + Style.RESET_ALL
    )

    final_code = lint_and_auto_fix(final_code, desc)

    optimize_program_runtime(final_code, desc)

    # Print a short summary for the user and append a log entry for this run
    print_option1_summary(desc)
    append_option1_log(desc)

    sys.exit(0)


def handle_option_2() -> None:
    """
    Option 2:
    ask for a GitHub repository url and a request,
    run gitingest and send digest plus request to Gemini Pro.
    """
    print(Fore.CYAN + "\n=== Option 2: fix or explain a GitHub repository ===" + Style.RESET_ALL)

    if ingest is None:
        print(
            Fore.RED
            + "\nOption 2 requires the 'gitingest' package. Please install it in your virtualenv with:"
            + Style.RESET_ALL
        )
        print("\n    pip install gitingest\n")
        print("Then run mega_coder.py again and choose option 2.")
        return

    repo_url = read_github_repo_url()
    if repo_url is None:
        return

    user_request = read_repo_request()
    if user_request is None:
        return

    try:
        repo_digest = build_repo_digest(repo_url)
    except RuntimeError as exc:
        print(
            Fore.RED
            + f"\n[gitingest] Could not create a digest for this repository: {exc}"
            + Style.RESET_ALL
        )
        return

    prompt = build_repo_prompt(repo_url, repo_digest, user_request)

    print(
        Fore.CYAN
        + "\n[gemini] Sending repository digest and your request to Gemini 2.5 Pro..."
        + Style.RESET_ALL
    )
    answer = generate_with_gemini_pro(prompt)

    print("\n=== Gemini 2.5 Pro answer ===")
    print(answer.strip())
    print("\n=== End of answer ===\n")


def handle_option_3() -> None:
    """
    Option 3:
    tell the user we will watch the screen and give realtime coding tips,
    then start the screen coaching loop.
    """
    print(
        Fore.CYAN
        + "\n=== Option 3: look at my screen and give realtime coding tips ==="
        + Style.RESET_ALL
    )

    print(
        Fore.GREEN
        + "Perfect. Show me your screen and I will be giving you tips on how to improve the code I see."
        + Style.RESET_ALL
    )

    start_screen_coaching_loop()


def main() -> None:
    """
    Entry point, show title and menu, handle user choices.
    """
    ensure_no_cli_args()

    while True:
        print()
        print(Fore.CYAN + APP_TITLE + Style.RESET_ALL)
        print()
        print(Fore.MAGENTA + MENU + Style.RESET_ALL)

        choice = read_menu_choice()
        if choice is None:
            break
        if choice == 1:
            handle_option_1()
        elif choice == 2:
            handle_option_2()
        elif choice == 3:
            handle_option_3()


if __name__ == "__main__":
    main()