import os
import json
import subprocess
import requests
import glob
import threading
import time
import signal
import tempfile
from enum import Enum
from git import Repo
import gradio as gr
from tenacity import retry, stop_after_attempt, wait_exponential

# ============================================================
# SECRETS — loaded from environment only, never hardcoded
# ============================================================
GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
GITHUB_REPO = os.getenv("GITHUB_REPO")
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")

# Enterprise Telegram Integration
TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
TELEGRAM_CHAT_ID = os.getenv("TELEGRAM_CHAT_ID")

# Fail fast at startup if critical orchestration secrets are missing
for _key, _val in [
    ("GITHUB_TOKEN", GITHUB_TOKEN),
    ("GITHUB_REPO", GITHUB_REPO),
    ("OPENROUTER_API_KEY", OPENROUTER_API_KEY),
]:
    if not _val:
        raise EnvironmentError(
            f"Required environment variable '{_key}' is not set. "
            "Set it in your HuggingFace Space secrets."
        )

# ============================================================
# PATHS & CONSTANTS
# ============================================================
PERSISTENT_DIR = "/data"
REPO_DIR = f"{PERSISTENT_DIR}/repo"
STATE_FILE = f"{PERSISTENT_DIR}/healing_state.json"
VENV_DIR = f"{PERSISTENT_DIR}/target_venv"
MCP_RUNTIME_CONFIG = "/tmp/mcp_runtime.json"
MODEL = "openrouter/qwen/qwen-2.5-coder-32b-instruct:free"

# ============================================================
# JOB STATE MACHINE
# ============================================================
class JobStatus(Enum):
    RUNNING = "RUNNING"
    DONE = "DONE"
    FAILED = "FAILED"

# ============================================================
# GLOBAL DASHBOARD STATE (thread-safe)
# ============================================================
dashboard_logs: list[str] = []
metrics = {"scanned": 0, "green": 0, "prs_created": 0, "status": "Idle"}
_audit_lock = threading.Lock()
_log_lock = threading.Lock()

def update_ui_log(message: str):
    timestamp = time.strftime("%H:%M:%S")
    formatted = f"[{timestamp}] {message}"
    print(formatted)
    with _log_lock:
        dashboard_logs.append(formatted)
        if len(dashboard_logs) > 100:
            dashboard_logs.pop(0)

# ============================================================
# ENTERPRISE TELEGRAM NOTIFICATIONS
# ============================================================
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
def _dispatch_telegram_payload(payload: dict):
    """Internal function wrapped in tenacity retry logic for resilient delivery."""
    url = f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendMessage"
    response = requests.post(url, json=payload, timeout=10)
    response.raise_for_status()

def send_telegram_alert(message: str):
    """
    Fire-and-forget notification system.
    Never blocks or crashes the main DevSecOps audit loop.
    """
    if not TELEGRAM_BOT_TOKEN or not TELEGRAM_CHAT_ID:
        return

    payload = {
        "chat_id": TELEGRAM_CHAT_ID,
        "text": message,
        "parse_mode": "Markdown",
        "disable_web_page_preview": True
    }
    
    try:
        _dispatch_telegram_payload(payload)
    except Exception as e:
        update_ui_log(f"⚠ Telegram alert failed after retries: {e}")

# ============================================================
# SECURE SUBPROCESS
# ============================================================
def run_subprocess_safe(
    cmd: list,
    cwd: str = REPO_DIR,
    timeout: int = 300,
    env_overrides: dict = None,
    raise_on_error: bool = True,
) -> tuple:
    if isinstance(cmd, str):
        raise TypeError("SECURITY VIOLATION: String commands are forbidden.")

    env = os.environ.copy()
    for secret_key in ["OPENROUTER_API_KEY", "GITHUB_TOKEN", "GITHUB_PERSONAL_ACCESS_TOKEN", "TELEGRAM_BOT_TOKEN"]:
        env.pop(secret_key, None)
    if env_overrides:
        env.update(env_overrides)

    proc = None
    try:
        proc = subprocess.Popen(
            cmd,
            shell=False,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            env=env,
            start_new_session=True,
        )
        stdout, stderr = proc.communicate(timeout=timeout)
        output = (stdout or "") + "\n" + (stderr or "")
        returncode = proc.returncode

        if raise_on_error and returncode != 0:
            raise subprocess.CalledProcessError(returncode, cmd, stdout, stderr)

        return output, returncode

    except subprocess.TimeoutExpired:
        if proc:
            try:
                os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
            except ProcessLookupError:
                pass
            proc.communicate()
        raise RuntimeError(f"Command timed out after {timeout}s: {cmd[0]}. Process group killed.")

# ============================================================
# SECRET-SAFE GIT CREDENTIALS
# ============================================================
def configure_git_credentials():
    cred_path = "/tmp/.git-credentials"
    with open(cred_path, "w") as f:
        f.write(f"https://x-token:{GITHUB_TOKEN}@github.com\n")
    os.chmod(cred_path, 0o600)

    run_subprocess_safe(
        ["git", "config", "--global", "credential.helper", f"store --file {cred_path}"],
        cwd="/tmp",
    )

# ============================================================
# MCP CONFIG
# ============================================================
def write_mcp_config() -> str:
    config = {
        "mcpServers": {
            "fetch-docs": {
                "command": "npx",
                "args": ["-y", "@modelcontextprotocol/server-fetch"],
                "env": {
                    "FETCH_ALLOWED_DOMAINS": "docs.python.org,pypi.org,docs.github.com,packaging.python.org,peps.python.org"
                },
            },
            "github-manager": {
                "command": "npx",
                "args": ["-y", "@modelcontextprotocol/server-github"],
                "env": {"GITHUB_PERSONAL_ACCESS_TOKEN": GITHUB_TOKEN},
            },
        }
    }

    with open(MCP_RUNTIME_CONFIG, "w") as f:
        json.dump(config, f, indent=2)
    os.chmod(MCP_RUNTIME_CONFIG, 0o600)
    return MCP_RUNTIME_CONFIG

# ============================================================
# ATOMIC STATE MACHINE
# ============================================================
def _load_raw_state() -> dict:
    if os.path.exists(STATE_FILE):
        try:
            with open(STATE_FILE, "r") as f:
                return json.load(f)
        except (json.JSONDecodeError, OSError):
            return {"jobs": {}}
    return {"jobs": {}}

def set_job_status(relative_test_path: str, status: JobStatus, detail: str = ""):
    with _audit_lock:
        state = _load_raw_state()
        state["jobs"][relative_test_path] = {
            "status": status.value,
            "detail": detail,
            "updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
        }
        tmp_path = STATE_FILE + ".tmp"
        with open(tmp_path, "w") as f:
            json.dump(state, f, indent=2)
        os.replace(tmp_path, STATE_FILE)

def get_job_status(relative_test_path: str):
    state = _load_raw_state()
    job = state.get("jobs", {}).get(relative_test_path)
    if not job:
        return None
    try:
        return JobStatus(job["status"])
    except ValueError:
        return None

# ============================================================
# VIRTUALENV ISOLATION FOR TARGET REPO (Powered by uv)
# ============================================================
def setup_target_venv() -> tuple:
    if not os.path.exists(VENV_DIR):
        update_ui_log("Creating isolated virtualenv using uv...")
        run_subprocess_safe(["uv", "venv", VENV_DIR])

    pytest_bin = os.path.join(VENV_DIR, "bin", "pytest")
    req_path = os.path.join(REPO_DIR, "requirements.txt")
    
    if os.path.exists(req_path):
        update_ui_log("Installing target repo deps via uv (blazing fast)...")
        run_subprocess_safe(
            ["uv", "pip", "install", "--python", VENV_DIR, "--quiet", "-r", req_path],
            cwd=REPO_DIR,
            timeout=600,
        )

    return None, pytest_bin

# ============================================================
# SAFE BRANCH MANAGEMENT
# ============================================================
def cleanup_stale_branch(branch_name: str):
    try:
        run_subprocess_safe(["git", "branch", "-D", branch_name], cwd=REPO_DIR, raise_on_error=False)
    except RuntimeError:
        pass
    try:
        run_subprocess_safe(["git", "push", "origin", "--delete", branch_name], cwd=REPO_DIR, raise_on_error=False)
    except RuntimeError:
        pass

def safe_git_pull():
    output, code = run_subprocess_safe(["git", "pull", "--ff-only", "origin", "main"], cwd=REPO_DIR, raise_on_error=False)
    if code != 0:
        update_ui_log(f"Pull non-fast-forward. Resetting hard to origin/main.")
        run_subprocess_safe(["git", "fetch", "origin"], cwd=REPO_DIR)
        run_subprocess_safe(["git", "reset", "--hard", "origin/main"], cwd=REPO_DIR)
        run_subprocess_safe(["git", "clean", "-fd"], cwd=REPO_DIR)

# ============================================================
# AIDER RUNNER
# ============================================================
def run_aider(mcp_config_path: str, prompt: str, context_files: list) -> tuple:
    prompt_fd, prompt_path = tempfile.mkstemp(prefix="aider_prompt_", suffix=".txt")
    try:
        with os.fdopen(prompt_fd, "w") as f:
            f.write(prompt)

        valid_context = [f for f in context_files if os.path.exists(os.path.join(REPO_DIR, f))]

        cmd = [
            "aider",
            "--model", MODEL,
            "--yes",
            "--no-stream",
            "--lint-cmd", "ruff check --fix .",
            "--mcp", mcp_config_path,
            "--message-file", prompt_path,
        ] + valid_context

        return run_subprocess_safe(
            cmd,
            cwd=REPO_DIR,
            timeout=600,
            env_overrides={"OPENROUTER_API_KEY": OPENROUTER_API_KEY},
            raise_on_error=False,
        )
    finally:
        try:
            os.unlink(prompt_path)
        except OSError:
            pass

# ============================================================
# MAIN AUDIT LOOP
# ============================================================
def enterprise_audit_loop():
    global metrics
    metrics["status"] = "Auditing Codebase..."
    update_ui_log("═" * 60)
    update_ui_log("Initiating Autonomous Healing Run.")
    send_telegram_alert("🚀 *Rhodawk DevSecOps*\n\nInitiating Autonomous Healing Run on target repository.")

    try:
        configure_git_credentials()
        mcp_config_path = write_mcp_config()

        if not os.path.exists(REPO_DIR):
            update_ui_log("Cloning repository (token via credential helper)...")
            Repo.clone_from(f"https://github.com/{GITHUB_REPO}.git", REPO_DIR)
            run_subprocess_safe(["git", "config", "user.name", "Rhodawk Auto-Healer"], cwd=REPO_DIR)
            run_subprocess_safe(["git", "config", "user.email", "security@rhodawk.ai"], cwd=REPO_DIR)
        else:
            update_ui_log("Syncing to latest origin/main...")
            safe_git_pull()

        _, pytest_bin = setup_target_venv()

        test_files = sorted(glob.glob(f"{REPO_DIR}/tests/test_*.py"))
        metrics["scanned"] = len(test_files)
        update_ui_log(f"Discovered {metrics['scanned']} test files.")

        for test_path in test_files:
            relative_test_path = os.path.relpath(test_path, REPO_DIR)
            filename = os.path.basename(relative_test_path)
            branch_name = f"security-patch/{filename.replace('.py', '').replace('_', '-')}"

            status = get_job_status(relative_test_path)

            if status == JobStatus.DONE:
                metrics["green"] += 1
                update_ui_log(f"⏭  Skipping (already DONE): {relative_test_path}")
                continue

            if status == JobStatus.RUNNING:
                update_ui_log(f"♻  Cleaning up interrupted job: {relative_test_path}")
                cleanup_stale_branch(branch_name)
                safe_git_pull()

            set_job_status(relative_test_path, JobStatus.RUNNING)
            update_ui_log(f"🔍 Testing: {relative_test_path}")

            logs, pytest_code = run_subprocess_safe(
                [pytest_bin, relative_test_path, "-v", "--tb=short"],
                cwd=REPO_DIR,
                timeout=120,
                raise_on_error=False,
            )

            if pytest_code == 0:
                update_ui_log(f"✅ PASSED: {relative_test_path}")
                set_job_status(relative_test_path, JobStatus.DONE, "tests passed")
                metrics["green"] += 1
                continue

            update_ui_log(f"❌ FAILED: {relative_test_path} — dispatching Aider agent...")
            send_telegram_alert(f"⚠️ *Test Failed*\n\nTarget: `{relative_test_path}`\nDispatching Aider agent for autonomous patch...")

            src_file = f"src/{filename.replace('test_', '')}"
            context_files = [relative_test_path]
            if os.path.exists(os.path.join(REPO_DIR, src_file)):
                context_files.append(src_file)
            if os.path.exists(os.path.join(REPO_DIR, "requirements.txt")):
                context_files.append("requirements.txt")

            prompt = (
                f"The pytest test '{relative_test_path}' failed with this output:\n\n"
                f"{logs}\n\n"
                f"Instructions:\n"
                f"1. If there is an import error or version conflict, use the "
                f"'fetch-docs' MCP tool to look up the correct library documentation.\n"
                f"2. Fix the source code in '{src_file}' or 'requirements.txt' to "
                f"make all tests pass. Do NOT modify test files.\n"
                f"3. Use the 'github-manager' MCP tool to create a new branch named "
                f"'{branch_name}' and open a Pull Request describing what was fixed.\n"
                f"4. Ensure the fix is minimal and does not introduce regressions."
            )

            aider_output, aider_code = run_aider(mcp_config_path, prompt, context_files)

            if aider_code != 0:
                update_ui_log(f"⚠  Aider non-zero exit for {relative_test_path}. Marking FAILED.")
                set_job_status(relative_test_path, JobStatus.FAILED, aider_output[-500:])
                send_telegram_alert(f"🔴 *Patch Generation Failed*\n\nTarget: `{relative_test_path}`\nAider exited with non-zero status.")
            else:
                update_ui_log(f"🔁 PR submitted for: {relative_test_path}")
                set_job_status(relative_test_path, JobStatus.DONE, "aider fix PR submitted")
                metrics["prs_created"] += 1
                send_telegram_alert(f"✅ *Auto-Heal PR Generated*\n\nTarget: `{relative_test_path}`\nAwaiting human review.")

            run_subprocess_safe(["git", "checkout", "main"], cwd=REPO_DIR, raise_on_error=False)
            safe_git_pull()

    except Exception as e:
        update_ui_log(f"🔴 FATAL: Audit loop crashed — {e}")
        send_telegram_alert(f"🔴 *FATAL ERROR*\n\nAudit loop crashed: `{e}`")
        metrics["status"] = "Error — see logs"
        return

    metrics["status"] = "Idle / Secure"
    update_ui_log("═" * 60)
    update_ui_log("Audit complete. All targets processed.")
    send_telegram_alert("🎉 *Audit Complete*\n\nAll targets neutralized or patched.")

# ============================================================
# GRADIO DASHBOARD
# ============================================================
def get_dashboard_data():
    with _log_lock:
        logs = "\n".join(dashboard_logs)
    return (
        logs,
        metrics["scanned"],
        metrics["green"],
        metrics["prs_created"],
        metrics["status"],
    )

def trigger_audit():
    with _audit_lock:
        if metrics["status"] not in ("Idle / Secure", "Idle"):
            return "Audit already running..."
        metrics["status"] = "Starting..."

    threading.Thread(target=enterprise_audit_loop, daemon=True).start()
    return "Audit Triggered!"

with gr.Blocks() as demo:
    gr.Markdown("# 🦅 Rhodawk AI — Autonomous DevSecOps Control Plane")
    gr.Markdown(
        "Agentic CI/CD: autonomous test healing with MCP-assisted PR generation. "
        "All changes require human review before merge."
    )

    with gr.Row():
        stat_status = gr.Textbox(label="System Status", value=metrics["status"], interactive=False)
        stat_scanned = gr.Number(label="Modules Scanned", value=metrics["scanned"], interactive=False)
        stat_green = gr.Number(label="Verified Green", value=metrics["green"], interactive=False)
        stat_prs = gr.Number(label="PRs Generated", value=metrics["prs_created"], interactive=False)

    with gr.Row():
        btn_trigger = gr.Button("🚀 Trigger Full System Audit", variant="primary")

    live_logs = gr.TextArea(
        label="Live Agent Execution Logs (MCP + Aider)",
        lines=20,
        interactive=False,
    )

    timer = gr.Timer(2)
    timer.tick(
        get_dashboard_data,
        inputs=None,
        outputs=[live_logs, stat_scanned, stat_green, stat_prs, stat_status]
    )

    demo.load(
        get_dashboard_data,
        inputs=None,
        outputs=[live_logs, stat_scanned, stat_green, stat_prs, stat_status]
    )
    btn_trigger.click(trigger_audit, inputs=None, outputs=stat_status)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Monochrome())