"""Background backend worker that owns research-job orchestration."""

from __future__ import annotations

import logging
import os
import threading
import time
from dataclasses import dataclass
from typing import Any

from backend.lean_checker import check_conjecture
from backend.research_heuristics import (
    build_initial_attempt_prompt,
    build_repair_prompt,
    generate_candidates,
)
from backend.supabase_runtime import get_supabase_client, supabase_is_configured
from backend.verified_lemma_store import store_verified_lemma
from postgrest.exceptions import APIError

logger = logging.getLogger(__name__)

ACTIVE_ATTEMPT_STATUSES = {"pending", "generated", "requeued"}
REPAIRABLE_STATUSES = {
    "typecheck_error",
    "timeout",
    "missing_imports",
    "malformed_output",
    "unsafe_placeholder",
    "checker_unavailable",
    "no_code",
    "failed",
    "expired",
}


@dataclass
class ResearchWorkerConfig:
    poll_interval_seconds: int = 5


class ResearchWorker:
    def __init__(self, config: ResearchWorkerConfig | None = None) -> None:
        self._config = config or ResearchWorkerConfig(
            poll_interval_seconds=max(1, int(os.environ.get("RESEARCH_BACKEND_POLL_INTERVAL", "5")))
        )
        self._stop_event = threading.Event()
        self._thread = threading.Thread(target=self._run, name="backend-research-loop", daemon=True)

    @classmethod
    def from_env(cls) -> "ResearchWorker | None":
        enabled = os.environ.get("RESEARCH_BACKEND_ENABLED", "0").strip().lower() not in {"0", "false", "no"}
        if not enabled or not supabase_is_configured():
            return None
        return cls()

    def start(self) -> None:
        self._thread.start()

    def stop(self) -> None:
        self._stop_event.set()
        self._thread.join(timeout=self._config.poll_interval_seconds + 2)

    def run_once(self) -> None:
        client = get_supabase_client()
        jobs_resp = client.rpc("backend_list_active_research_jobs", {}).execute()
        jobs = jobs_resp.data or []
        for job in jobs:
            job_id = job.get("id")
            if not job_id:
                continue
            try:
                self._process_job(job_id)
            except Exception as exc:  # pragma: no cover - exercised in integration
                logger.exception("Research job %s failed inside backend loop", job_id)
                client.rpc(
                    "backend_update_research_job_status",
                    {
                        "p_job_id": job_id,
                        "p_status": "failed",
                        "p_latest_error": str(exc),
                    },
                ).execute()

    def _run(self) -> None:
        while not self._stop_event.is_set():
            self.run_once()
            self._stop_event.wait(self._config.poll_interval_seconds)

    def _process_job(self, job_id: str) -> None:
        client = get_supabase_client()
        state_resp = client.rpc("backend_get_research_job_state", {"p_job_id": job_id}).execute()
        job = state_resp.data or {}
        if not job or job.get("status") not in {"active", "pending"}:
            return

        self._verify_generated_attempts(job)
        state_resp = client.rpc("backend_get_research_job_state", {"p_job_id": job_id}).execute()
        job = state_resp.data or job

        episodes = list(job.get("episodes", []))
        for episode in episodes:
            self._process_episode(job, episode)

        state_resp = client.rpc("backend_get_research_job_state", {"p_job_id": job_id}).execute()
        job = state_resp.data or job
        enqueued_new_work = self._ensure_capacity(job)
        self._maybe_finish_job(job_id, enqueued_new_work=enqueued_new_work)

    def _verify_generated_attempts(self, job: dict[str, Any]) -> None:
        client = get_supabase_client()
        for episode in job.get("episodes", []):
            for attempt in episode.get("attempts", []):
                if attempt.get("verification_status") != "generated":
                    continue
                result = check_conjecture(attempt.get("generated_lean_code") or attempt.get("raw_output", ""))
                verifier_metadata: dict[str, Any] = {
                    "compiler": result.compiler,
                    "is_true": result.is_true,
                }
                if result.is_true:
                    try:
                        verifier_metadata["document_id"] = store_verified_lemma(
                            {
                                "source_task_id": attempt.get("task_id"),
                                "source_prompt": attempt.get("prompt"),
                                "metadata": {
                                    "job_id": job.get("id"),
                                    "episode_id": episode.get("id"),
                                    "attempt_id": attempt.get("id"),
                                    "domain": episode.get("domain"),
                                },
                                "conjecture": result.conjecture,
                                "normalized_code": result.normalized_code,
                                "verification_status": result.verification_status,
                                "verification_summary": result.verification_summary,
                                "compiler": result.compiler,
                                "stdout": result.stdout,
                                "stderr": result.stderr,
                                "is_true": result.is_true,
                            }
                        )
                    except Exception as exc:  # pragma: no cover - external dependency
                        verifier_metadata["storage_error"] = str(exc)
                client.rpc(
                    "backend_verify_research_attempt",
                    {
                        "p_attempt_id": attempt["id"],
                        "p_verification_status": result.verification_status,
                        "p_verification_summary": result.verification_summary,
                        "p_normalized_code": result.normalized_code,
                        "p_lean_stdout": result.stdout,
                        "p_lean_stderr": result.stderr,
                        "p_verifier_metadata": verifier_metadata,
                    },
                ).execute()

    def _process_episode(self, job: dict[str, Any], episode: dict[str, Any]) -> None:
        if episode.get("status") != "active":
            return

        attempts = list(episode.get("attempts", []))
        if any(attempt.get("verification_status") in ACTIVE_ATTEMPT_STATUSES for attempt in attempts):
            return

        if int(episode.get("attempt_count", 0) or 0) >= int(episode.get("max_attempts", 0) or 0):
            return

        latest_attempt = attempts[-1] if attempts else None
        if latest_attempt and latest_attempt.get("verification_status") not in REPAIRABLE_STATUSES:
            return

        if self._remaining_budget(job) < int(job.get("budget_per_attempt", 0) or 0):
            return

        prompt = build_repair_prompt(episode, latest_attempt or {})
        metadata = {
            "domain": episode.get("domain", ""),
            "conjecture_family": episode.get("strategy_tag", ""),
            "interestingness_score": episode.get("interestingness_score", 0.0),
            "novelty_score": episode.get("novelty_score", 0.0),
            "repair_objective": "Fix Lean compilation errors and prove the conjecture without placeholders.",
            "previous_code": (latest_attempt or {}).get("generated_lean_code", ""),
            "lean_stderr": (latest_attempt or {}).get("lean_stderr", ""),
            "lean_stdout": (latest_attempt or {}).get("lean_stdout", ""),
            "proof_hint": (episode.get("metadata", {}) or {}).get("proof_hint", ""),
        }
        try:
            get_supabase_client().rpc(
                "backend_enqueue_research_attempt",
                {
                    "p_job_id": job["id"],
                    "p_episode_id": episode["id"],
                    "p_prompt": prompt,
                    "p_model": job.get("model", ""),
                    "p_task_type": "proof_repair",
                    "p_strategy_tag": "repair",
                    "p_parent_attempt_id": (latest_attempt or {}).get("id"),
                    "p_conjecture_text": episode.get("conjecture_text", ""),
                    "p_formal_statement": episode.get("formal_statement", ""),
                    "p_metadata": metadata,
                },
            ).execute()
        except APIError as exc:
            if _is_attempt_budget_exhausted_error(exc):
                logger.info(
                    "Skipping repair enqueue for exhausted episode %s",
                    episode.get("id", "<unknown>"),
                )
                return
            raise

    def _ensure_capacity(self, job: dict[str, Any]) -> bool:
        budget_per_attempt = int(job.get("budget_per_attempt", 0) or 0)
        if budget_per_attempt <= 0 or self._remaining_budget(job) < budget_per_attempt:
            return False

        active_episode_count = 0
        verified_episodes: list[dict[str, Any]] = []
        all_episodes = list(job.get("episodes", []))
        for episode in all_episodes:
            if episode.get("status") == "verified":
                verified_episodes.append(episode)
            if episode.get("status") == "active":
                active_episode_count += 1

        max_concurrent = int(job.get("max_concurrent_episodes", 1) or 1)
        if active_episode_count >= max_concurrent:
            return False

        candidates = generate_candidates(
            verified_episodes=verified_episodes,
            desired_domain=str(job.get("domain", "") or ""),
            seed_conjecture=str(job.get("seed_conjecture", "") or "") if not all_episodes else "",
        )
        existing_conjectures = {
            str(episode.get("conjecture_text", "")).strip().lower()
            for episode in all_episodes
        }
        for candidate in candidates:
            if candidate.conjecture_text.strip().lower() in existing_conjectures:
                continue
            client = get_supabase_client()
            episode_resp = client.rpc(
                "backend_create_research_episode",
                {
                    "p_job_id": job["id"],
                    "p_conjecture_text": candidate.conjecture_text,
                    "p_formal_statement": candidate.formal_statement,
                    "p_domain": candidate.domain,
                    "p_strategy_tag": candidate.strategy_tag,
                    "p_metadata": candidate.metadata,
                    "p_interestingness_score": candidate.interestingness_score,
                    "p_novelty_score": candidate.novelty_score,
                },
            ).execute()
            episode_id = episode_resp.data
            if not episode_id:
                return False
            client.rpc(
                "backend_enqueue_research_attempt",
                {
                    "p_job_id": job["id"],
                    "p_episode_id": episode_id,
                    "p_prompt": build_initial_attempt_prompt(candidate),
                    "p_model": job.get("model", ""),
                    "p_task_type": "proof_attempt",
                    "p_strategy_tag": candidate.strategy_tag,
                    "p_parent_attempt_id": None,
                    "p_conjecture_text": candidate.conjecture_text,
                    "p_formal_statement": candidate.formal_statement,
                    "p_metadata": {
                        "domain": candidate.domain,
                        "conjecture_family": candidate.strategy_tag,
                        "interestingness_score": candidate.interestingness_score,
                        "novelty_score": candidate.novelty_score,
                        "proof_hint": candidate.metadata.get("proof_hint", ""),
                    },
                },
            ).execute()
            return True
        return False

    def _maybe_finish_job(self, job_id: str, enqueued_new_work: bool = False) -> None:
        client = get_supabase_client()
        state_resp = client.rpc("backend_get_research_job_state", {"p_job_id": job_id}).execute()
        job = state_resp.data or {}
        if not job:
            return
        episodes = list(job.get("episodes", []))
        has_active_attempt = any(
            attempt.get("verification_status") in ACTIVE_ATTEMPT_STATUSES
            for episode in episodes
            for attempt in episode.get("attempts", [])
        )
        if has_active_attempt:
            return

        if enqueued_new_work:
            return

        active_episodes = [episode for episode in episodes if episode.get("status") == "active"]
        remaining_budget = self._remaining_budget(job)
        budget_per_attempt = int(job.get("budget_per_attempt", 0) or 0)
        if active_episodes and remaining_budget >= budget_per_attempt:
            return

        if remaining_budget >= budget_per_attempt and episodes:
            pass
        elif remaining_budget >= budget_per_attempt and not episodes:
            return

        status = "completed" if any(episode.get("status") == "verified" for episode in episodes) else "exhausted"
        client.rpc(
            "backend_update_research_job_status",
            {
                "p_job_id": job_id,
                "p_status": status,
                "p_latest_error": "",
            },
        ).execute()

    @staticmethod
    def _remaining_budget(job: dict[str, Any]) -> int:
        total_budget = int(job.get("total_budget", 0) or 0)
        credits_committed = int(job.get("credits_committed", 0) or 0)
        return max(total_budget - credits_committed, 0)


def _is_attempt_budget_exhausted_error(exc: APIError) -> bool:
    return "exhausted its attempt budget" in str(exc).lower()
