"""
Researcher — Autonomous research mining for LITCOIN.

    from litcoin import Agent

    agent = Agent(bankr_key="bk_...", ai_key="sk-...", ai_url="https://...")

    # Single research cycle
    result = agent.research_mine(task_type="code_optimization")

    # Autonomous loop
    agent.research_loop(task_type="code_optimization", rounds=10)
"""

import time
import json
import logging
import subprocess
import sys
import tempfile
import os

from litcoin.api import CoordinatorAPI, APIError

log = logging.getLogger("litcoin")

MAX_LOCAL_RUNTIME = 120  # seconds

# ─── Security: blocked patterns in LLM-generated code ──────────────────────
# Matches coordinator server-side sanitization (research.js sanitizeCode)
BLOCKED_PATTERNS = [
    "import os", "import sys", "import subprocess", "import shutil",
    "os.system", "os.popen", "subprocess.", "exec(", "eval(",
    "__import__", "open(", "pathlib", "importlib",
    "import socket", "import http", "import urllib", "import requests",
    "import ctypes", "import signal", "import multiprocessing",
    "from os ", "from sys ", "from subprocess ", "from shutil ",
    "from socket ", "from http ", "from pathlib ", "from ctypes ",
]

_SECURITY_WARNING_SHOWN = False


def _show_security_warning():
    """Show one-time security warning about local code execution."""
    global _SECURITY_WARNING_SHOWN
    if _SECURITY_WARNING_SHOWN:
        return
    _SECURITY_WARNING_SHOWN = True
    log.warning("=" * 60)
    log.warning("SECURITY NOTICE: Research mining generates and executes")
    log.warning("Python code locally via LLM. Code is sanitized to block")
    log.warning("dangerous imports (os, subprocess, socket, etc.) but")
    log.warning("runs in a subprocess with limited isolation.")
    log.warning("Review: https://litcoiin.xyz/docs#security")
    log.warning("=" * 60)


def sanitize_code(code):
    """
    Check generated code for dangerous patterns.
    Returns (safe: bool, reason: str or None).
    Mirrors coordinator-side sanitizeCode() in research.js.
    """
    for pattern in BLOCKED_PATTERNS:
        if pattern in code:
            return False, f"Blocked pattern: {pattern}"
    return True, None


def extract_reasoning(raw_response):
    """
    Extract reasoning trace from LLM response.
    Handles <think> tags (DeepSeek-R1, QwQ), prose before code blocks, etc.
    Returns (reasoning: str or None, remaining_content: str).
    """
    import re

    reasoning = None
    content = raw_response

    # DeepSeek-R1 / QwQ style: <think>...</think>
    think_match = re.search(r'<think>(.*?)</think>', raw_response, re.DOTALL)
    if think_match:
        reasoning = think_match.group(1).strip()
        content = raw_response[think_match.end():].strip()
        return reasoning[:50000] if reasoning else None, content

    # Prose before code block (at least 50 chars of thinking before ```)
    if '```' in raw_response:
        idx = raw_response.index('```')
        before = raw_response[:idx].strip()
        if len(before) > 50:
            reasoning = before
            content = raw_response[idx:]

    return (reasoning[:50000] if reasoning else None), content


class Researcher:
    """
    Handles the research mining loop:
    1. Fetch a task from the coordinator
    2. Use LLM to generate/iterate on experiment code
    3. Run experiment locally (sandboxed)
    4. Submit result if it improves on baseline
    """

    def __init__(self, api, wallet, ai_config=None):
        """
        Args:
            api: CoordinatorAPI instance
            wallet: Miner wallet address
            ai_config: dict with keys: key, url, model, anthropic_mode
        """
        self.api = api
        self.wallet = wallet
        self.ai = ai_config or {}
        self._history = {}  # taskId -> list of (code, metric) attempts

    def list_tasks(self, task_type=None):
        """List available research tasks."""
        return self.api.research_list_tasks(task_type)

    def get_task(self, task_id):
        """Get full task details."""
        return self.api.research_get_task(task_id)

    def fetch_task(self, task_type=None):
        """Fetch a random task for mining."""
        return self.api.research_fetch_task(task_type, self.wallet)

    def submit(self, task_id, code, reasoning=None):
        """Submit code to coordinator for verification."""
        model = self.ai.get("model")
        provider = self.ai.get("url", "").split("//")[-1].split("/")[0] if self.ai.get("url") else None
        return self.api.research_submit(task_id, self.wallet, code, model=model, model_provider=provider, reasoning=reasoning)

    def leaderboard(self, task_id=None):
        """Get research leaderboard."""
        return self.api.research_leaderboard(task_id)

    def stats(self):
        """Get global research stats."""
        return self.api.research_stats()

    def results(self, task_id=None, limit=20):
        """Get verified research results (discoveries)."""
        return self.api.research_results(task_id=task_id, miner=self.wallet, limit=limit)

    # ─── Local Experiment Runner ──────────────────────────────────────────

    def run_locally(self, code, test_code, entry_function, timeout=MAX_LOCAL_RUNTIME):
        """
        Run experiment code locally in a sandboxed subprocess.
        
        Security measures:
        - Code is sanitized against dangerous imports (os, subprocess, socket, etc.)
        - Subprocess runs with minimal environment variables (no HOME, no shell config)
        - Working directory is a temp folder (no access to user files)
        - Timeout enforced to prevent infinite loops
        - Temp script file is deleted after execution
        
        Returns dict with: success, metric_name, metric_value, stdout, stderr, error
        """
        # ── Security: sanitize LLM-generated code before execution ──
        safe, reason = sanitize_code(code)
        if not safe:
            return {"success": False, "error": f"Code rejected (security): {reason}"}

        # Also check test_code in case the task harness is compromised
        safe_t, reason_t = sanitize_code(test_code)
        if not safe_t:
            return {"success": False, "error": f"Test harness rejected (security): {reason_t}"}

        full_script = f"""{code}

# ─── Test harness ───
try:
{chr(10).join('    ' + line for line in test_code.strip().split(chr(10)))}
    print("VERIFY:PASS")
except AssertionError as e:
    print(f"VERIFY:FAIL:{{e}}")
except Exception as e:
    print(f"VERIFY:ERROR:{{e}}")
"""

        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
            f.write(full_script)
            script_path = f.name

        try:
            # ── Restricted environment: minimal PATH, no HOME/USER/SHELL ──
            safe_env = {
                "PATH": os.path.dirname(sys.executable),
                "PYTHONPATH": "",
                "PYTHONDONTWRITEBYTECODE": "1",
                "TMPDIR": tempfile.gettempdir(),
                "LANG": "en_US.UTF-8",
            }

            result = subprocess.run(
                [sys.executable, "-I", script_path],  # -I = isolated mode (no user site-packages, no PYTHON* env vars)
                capture_output=True, text=True,
                timeout=timeout,
                cwd=tempfile.gettempdir(),
                env=safe_env,
            )

            stdout = result.stdout
            stderr = result.stderr

            if result.returncode != 0:
                return {
                    "success": False,
                    "error": f"Exit code {result.returncode}: {stderr[:500]}",
                    "stdout": stdout[:1000],
                    "stderr": stderr[:1000],
                }

            if "VERIFY:FAIL" in stdout:
                fail_msg = ""
                for line in stdout.split("\n"):
                    if "VERIFY:FAIL" in line:
                        fail_msg = line.split("VERIFY:FAIL:")[-1]
                return {"success": False, "error": f"Correctness failed: {fail_msg}"}

            if "VERIFY:ERROR" in stdout:
                err_msg = ""
                for line in stdout.split("\n"):
                    if "VERIFY:ERROR" in line:
                        err_msg = line.split("VERIFY:ERROR:")[-1]
                return {"success": False, "error": f"Runtime error: {err_msg}"}

            if "VERIFY:PASS" not in stdout:
                return {"success": False, "error": "Verification did not pass"}

            # Extract metric
            metric_name = None
            metric_value = None
            for line in stdout.split("\n"):
                if line.startswith("METRIC:"):
                    parts = line.split(":")
                    if len(parts) >= 3:
                        metric_name = parts[1]
                        try:
                            metric_value = float(parts[2])
                        except ValueError:
                            pass

            if metric_value is None:
                return {"success": False, "error": "No metric reported"}

            return {
                "success": True,
                "metric_name": metric_name,
                "metric_value": metric_value,
                "stdout": stdout[:2000],
            }

        except subprocess.TimeoutExpired:
            return {"success": False, "error": f"Timeout ({timeout}s)"}
        except Exception as e:
            return {"success": False, "error": str(e)}
        finally:
            try:
                os.unlink(script_path)
            except:
                pass

    # ─── LLM Code Generation ─────────────────────────────────────────────

    def _call_llm(self, prompt, max_tokens=4096):
        """Call the miner's LLM to generate experiment code."""
        if not self.ai.get("key"):
            raise RuntimeError("AI key required for research mining. Set ai_key when creating Agent.")

        import requests

        headers = {"Content-Type": "application/json"}
        url = self.ai["url"]
        model = self.ai.get("model", "gpt-4o-mini")

        if self.ai.get("anthropic_mode"):
            headers["x-api-key"] = self.ai["key"]
            headers["anthropic-version"] = "2023-06-01"
            payload = {
                "model": model,
                "max_tokens": max_tokens,
                "messages": [{"role": "user", "content": prompt}],
            }
        else:
            headers["Authorization"] = f"Bearer {self.ai['key']}"
            payload = {
                "model": model,
                "max_tokens": max_tokens,
                "messages": [{"role": "user", "content": prompt}],
                "temperature": 0.7,
            }

        r = requests.post(
            url if "/v1/" in url else f"{url}/chat/completions",
            headers=headers,
            json=payload,
            timeout=120,
        )
        r.raise_for_status()
        data = r.json()

        if self.ai.get("anthropic_mode"):
            return data["content"][0]["text"]
        return data["choices"][0]["message"]["content"]

    def _extract_code(self, llm_response):
        """Extract Python code from LLM response (handles markdown fences)."""
        # Try to find ```python ... ``` block
        if "```python" in llm_response:
            start = llm_response.index("```python") + 9
            end = llm_response.index("```", start)
            return llm_response[start:end].strip()
        elif "```" in llm_response:
            start = llm_response.index("```") + 3
            # Skip language identifier on first line
            if "\n" in llm_response[start:start+20]:
                start = llm_response.index("\n", start) + 1
            end = llm_response.index("```", start)
            return llm_response[start:end].strip()
        return llm_response.strip()

    def generate_solution(self, task, previous_attempts=None):
        """
        Use LLM to generate a solution for a research task.

        Args:
            task: Task dict from coordinator
            previous_attempts: List of (code, metric_value) tuples

        Returns:
            tuple: (code: str, reasoning: str or None)
        """
        prompt = f"""You are a research engineer competing in an optimization challenge.

TASK: {task['title']}
TYPE: {task['type']}
DESCRIPTION: {task['description']}

CONSTRAINTS:
- Language: {task['constraints']['language']}
- Entry function: `{task['constraints']['entry_function']}`
- Max runtime: {task['constraints']['max_runtime']} seconds

BASELINE:
- Metric: {task['baseline']['metric']}
- Baseline value: {task['baseline']['value']}
- Direction: {task['baseline']['direction']}
- Current best: {task.get('bestResult', {}).get('value', 'none yet') if task.get('bestResult') else 'none yet'}

"""

        if previous_attempts:
            prompt += "PREVIOUS ATTEMPTS (improve on these):\n"
            for i, (prev_code, prev_metric) in enumerate(previous_attempts[-3:]):
                prompt += f"\nAttempt {i+1} — metric: {prev_metric}\n```python\n{prev_code}\n```\n"
            prompt += "\nAnalyze what worked and what didn't. Try a fundamentally different approach.\n"
        else:
            prompt += "This is your first attempt. Write the best solution you can.\n"

        prompt += f"""
Think through your approach step by step before writing code.
Then write ONLY the Python function `{task['constraints']['entry_function']}` (and any helpers it needs).
Do NOT include test code, imports of os/sys/subprocess, or anything dangerous.
Focus on correctness first, then optimization."""

        response = self._call_llm(prompt)
        reasoning, remaining = extract_reasoning(response)
        code = self._extract_code(remaining)
        return code, reasoning

    # ─── Research Mining Cycle ────────────────────────────────────────────

    def mine_once(self, task_type=None, task_id=None):
        """
        Run one research mining cycle:
        1. Fetch task
        2. Generate solution via LLM
        3. Test locally
        4. Submit if it beats baseline

        Returns dict with: task, code, local_result, submission (if submitted)
        """
        # Show one-time security warning about local code execution
        _show_security_warning()

        # 1. Get task
        if task_id:
            task = self.get_task(task_id)
        else:
            task = self.fetch_task(task_type)

        log.info(f"[research] Task: {task['title']} (baseline: {task['baseline']['value']})")

        # 2. Get previous attempts for this task
        prev = self._history.get(task["id"], [])

        # 3. Generate code
        code, reasoning = self.generate_solution(task, prev if prev else None)
        log.info(f"[research] Generated {len(code)} chars of code{f' + {len(reasoning)} chars reasoning' if reasoning else ''}")

        # 4. Test locally
        local = self.run_locally(
            code,
            task["constraints"].get("test_code", ""),
            task["constraints"]["entry_function"],
            timeout=task["constraints"].get("max_runtime", MAX_LOCAL_RUNTIME),
        )

        result = {
            "task": {"id": task["id"], "title": task["title"]},
            "code": code,
            "reasoning": reasoning,
            "local_result": local,
            "submission": None,
        }

        if not local["success"]:
            log.warning(f"[research] Local test failed: {local['error']}")
            # Track failed attempt so LLM can learn from it
            self._history.setdefault(task["id"], []).append((code, None))
            return result

        metric = local["metric_value"]
        log.info(f"[research] Local metric: {local['metric_name']} = {metric}")

        # Track attempt
        self._history.setdefault(task["id"], []).append((code, metric))

        # 5. Check if it beats baseline before submitting
        baseline = task["baseline"]["value"]
        direction = task["baseline"]["direction"]
        beats = (metric < baseline) if direction == "lower_is_better" else (metric > baseline)

        if not beats:
            log.info(f"[research] Did not beat baseline ({baseline}), skipping submit")
            return result

        # 6. Submit to coordinator for full verification
        log.info(f"[research] Submitting — beats baseline ({metric} vs {baseline})")
        try:
            sub = self.submit(task["id"], code, reasoning=reasoning)
            result["submission"] = sub
            if sub.get("verified"):
                log.info(f"[research] ✓ Verified! {sub.get('message', '')}")
                if sub.get("isNewBest"):
                    log.info(f"[research] 🏆 NEW RECORD!")
            else:
                log.warning(f"[research] ✗ Verification failed: {sub.get('error', 'unknown')}")
        except Exception as e:
            log.error(f"[research] Submit error: {e}")

        return result

    def mine_loop(self, task_type=None, task_id=None, rounds=10, delay=30):
        """
        Run autonomous research mining loop (Karpathy-style iteration).

        Sticks to ONE task and iterates on it. Each round builds on
        previous attempts' learnings. This is where real breakthroughs happen.

        Args:
            task_type: Filter by type (optional)
            task_id: Specific task to iterate on (recommended)
            rounds: Number of iterations (default 10)
            delay: Seconds between iterations (default 30)
        """
        log.info(f"[research] Starting research loop: {rounds} iterations, {'task ' + task_id if task_id else task_type or 'any'}")

        current_task_id = task_id

        for i in range(rounds):
            log.info(f"\n[research] ─── Iteration {i+1}/{rounds} ───")
            try:
                result = self.mine_once(task_type=task_type, task_id=current_task_id)

                # Lock onto this task for subsequent rounds
                if not current_task_id and result.get("task"):
                    current_task_id = result["task"]["id"]
                    log.info(f"[research] Locked onto: {result['task']['title']}")

                if result["submission"] and result["submission"].get("verified"):
                    reward = result["submission"].get("reward", 0)
                    streak = result["submission"].get("streak", 0)
                    pb = result["submission"].get("isPersonalBest", False)
                    nb = result["submission"].get("isNewBest", False)
                    marker = "🏆 GLOBAL RECORD" if nb else "⬆ Personal best" if pb else "✓"
                    log.info(f"[research] {marker} — iteration #{result['submission'].get('iteration', i+1)}, streak {streak}, +{reward:,} LITCOIN")
                elif result["local_result"]["success"]:
                    log.info(f"[research] Iteration {i+1}: metric={result['local_result']['metric_value']}")
                else:
                    log.info(f"[research] Iteration {i+1}: local test failed")

            except Exception as e:
                log.error(f"[research] Iteration {i+1} error: {e}")

            if i < rounds - 1:
                time.sleep(delay)

        log.info(f"[research] Loop complete. {rounds} iterations on {'task ' + current_task_id if current_task_id else 'various tasks'}.")
