"""
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


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):
        """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)

    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 subprocess.
        Returns dict with: success, metric_name, metric_value, stdout, stderr, error
        """
        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:
            result = subprocess.run(
                [sys.executable, script_path],
                capture_output=True, text=True,
                timeout=timeout,
                cwd=tempfile.gettempdir(),
            )

            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:
            str: Python code for the solution
        """
        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"""
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.
Respond with ONLY the Python code, no explanation."""

        response = self._call_llm(prompt)
        return self._extract_code(response)

    # ─── 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)
        """
        # 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 = self.generate_solution(task, prev if prev else None)
        log.info(f"[research] Generated {len(code)} chars of code")

        # 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,
            "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)
            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'}.")
