"""
model_io.py — Load the base + LoRA adapter once, expose two functions:

  score_pair(team_ctx, paper) → float
     Computes logprob(paper | team_ctx) with the LoRA adapter ENABLED.
     This is the preference signal the v1.4 LoRA was trained for.

  generate_pr_body(team_ctx, paper) → str
     Generates a PR-body draft with the LoRA adapter DISABLED.
     We use the base model for generation because the LoRA was trained
     for preference distillation, not free-text generation — empirically
     it emits EOS immediately if asked to generate.

Both functions assume the model is on GPU and toggle adapters in/out
inside the @spaces.GPU function wrapper.
"""
from __future__ import annotations

import os
from typing import Optional

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# TODO: swap to v1.4 once training completes.
BASE_MODEL = "Qwen/Qwen3.5-2B"
LORA_ADAPTER = os.environ.get(
    "MHPD_ADAPTER", "remyxai/mhpd-dpo-qwen3.5-2b-lora-v1.2-ipo-2epoch"
)

# "local"  → use the base 2B (LoRA off) for PR-body generation, inside @spaces.GPU
# "gemini" → use Gemini Flash via llm_api, on CPU (no GPU contention)
#
# Tradeoff: "local" keeps everything self-hosted but adds 3-6s of GPU time
# per request. "gemini" is faster (~2-3s, CPU-only) and frees Zero-GPU quota
# for the scoring step, at ~$0.0001-0.001 per request.
GENERATION_BACKEND = os.environ.get("MHPD_GENERATION_BACKEND", "local").lower()

# Globals — populated by load() at startup
_model: Optional[AutoModelForCausalLM] = None
_tokenizer: Optional[AutoTokenizer] = None


def load() -> tuple[AutoModelForCausalLM, AutoTokenizer]:
    """Load base + LoRA. Called once at Space cold-start.

    Uses transformers' PEFT integration (`base.load_adapter(...)`) rather
    than PeftModel.from_pretrained. The latter wraps the base in a
    PeftModel, but transformers 5.x's enable_adapters/disable_adapters
    integration hooks operate on the base model — and looking through a
    PeftModel wrapper they see "no adapter loaded" and raise. The
    load_adapter API loads the adapter directly onto the base model
    where the integration expects to find it.
    """
    global _model, _tokenizer
    if _model is not None:
        return _model, _tokenizer

    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
    # Make sure BOTH env vars are populated so anything that auto-reads
    # them (load_adapter, hf_hub_download via transformers' integration)
    # picks up auth.
    if hf_token:
        os.environ["HF_TOKEN"] = hf_token
        os.environ["HUGGINGFACE_TOKEN"] = hf_token
    print(f"[model_io] Loading {BASE_MODEL} + adapter {LORA_ADAPTER}…")
    _tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True, token=hf_token)
    if _tokenizer.pad_token is None:
        _tokenizer.pad_token = _tokenizer.eos_token

    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        dtype=torch.bfloat16,
        device_map="auto",
        attn_implementation="sdpa",
        trust_remote_code=True,
        token=hf_token,
    )
    base.config.use_cache = True   # generation needs it; scoring doesn't care

    # Load LoRA via transformers' PEFT integration.
    # Note: this transformers build's load_adapter doesn't accept a
    # `token` kwarg — it reads HF_TOKEN/HUGGINGFACE_TOKEN from env.
    # We set both at startup just above, so no need to pass it here.
    base.load_adapter(LORA_ADAPTER, adapter_name="default")
    # Explicitly mark the adapter active. Some transformers builds load
    # into peft_config but don't activate, leading to "No adapter loaded"
    # errors on later enable/disable calls.
    try:
        base.set_adapter("default")
    except Exception as e:
        print(f"[model_io] set_adapter warning: {e}")
    # Diagnostics — adapter state visible in container logs at startup,
    # without firing a request.
    pc = getattr(base, "peft_config", {}) or {}
    print(f"[model_io] peft_config keys: {list(pc.keys())}")
    try:
        active = base.active_adapters()
        print(f"[model_io] active adapters: {active}")
    except Exception as e:
        print(f"[model_io] active_adapters() warning: {e}")
    base.eval()
    _model = base
    print("[model_io] Ready.")
    return _model, _tokenizer


def score_pair(team_ctx: str, paper_text: str) -> float:
    """logprob(paper | team_ctx) — the team-aligned preference score.

    Higher score = the LoRA-adapted model expects this paper as a more-
    natural continuation given the team's history. This is exactly the
    quantity v1.4 was trained to model on chosen vs rejected pairs.

    Call inside @spaces.GPU with adapters enabled.
    """
    assert _model is not None, "Call load() first"
    messages = [
        {"role": "user", "content": team_ctx},
        {"role": "assistant", "content": paper_text},
    ]
    full_text = _tokenizer.apply_chat_template(messages, tokenize=False)

    prompt_msg = [{"role": "user", "content": team_ctx}]
    prompt_text = _tokenizer.apply_chat_template(
        prompt_msg, tokenize=False, add_generation_prompt=True
    )

    inputs = _tokenizer(full_text, return_tensors="pt", truncation=True, max_length=4096).to(_model.device)
    prompt_len = _tokenizer(prompt_text, return_tensors="pt").input_ids.shape[1]

    with torch.no_grad():
        out = _model(**inputs)
    log_probs = torch.log_softmax(out.logits[:, :-1, :], dim=-1)
    target_ids = inputs.input_ids[:, 1:]
    token_lp = log_probs.gather(2, target_ids.unsqueeze(-1)).squeeze(-1)
    # Average log-prob of the assistant turn (the paper text)
    completion_lp = token_lp[:, prompt_len - 1 :]
    return float(completion_lp.mean().item())


def generate_pr_body(team_ctx: str, paper_text: str, max_new_tokens: int = 350) -> str:
    """Router — dispatches to local 2B or Gemini Flash per GENERATION_BACKEND.

    Local path runs inside @spaces.GPU (with adapters disabled). Gemini
    path runs on CPU — call this OUTSIDE the GPU function to avoid
    burning GPU time on what is effectively an HTTP request.
    """
    if GENERATION_BACKEND == "gemini":
        # Lazy import so the space starts without google-generativeai
        # if the user has GENERATION_BACKEND=local and didn't install it.
        from llm_api import gemini_generate_pr_body
        return gemini_generate_pr_body(team_ctx, paper_text)
    return _generate_pr_body_local(team_ctx, paper_text, max_new_tokens)


def generation_runs_on_gpu() -> bool:
    """Orchestration helper. True iff the current backend needs the GPU
    container for PR body generation. The app uses this to decide
    whether to fold generation into the @spaces.GPU function or call it
    on CPU after scoring returns."""
    return GENERATION_BACKEND == "local"


def _generate_pr_body_local(team_ctx: str, paper_text: str, max_new_tokens: int = 600) -> str:
    """Draft a PR body integrating `paper_text` into the team's codebase.

    Runs with adapters DISABLED (base model). The v1.2/v1.4 LoRA was
    trained for ranking, not generation — empirically it emits EOS on
    first token if asked to generate. Disabling the adapter restores
    the base model's instruction-following capability.

    Caller is responsible for being inside @spaces.GPU and toggling
    adapters off before calling this.
    """
    assert _model is not None, "Call load() first"
    user_prompt = (
        "You are a software engineer integrating a research paper into your "
        "team's codebase. Write a concise pull request body proposing the "
        "integration.\n\n"
        f"Team context:\n{team_ctx[:1500]}\n\n"
        f"Paper:\n{paper_text[:1800]}\n\n"
        "Write the PR body now using this structure:\n\n"
        "## Summary\n(1-2 sentences)\n\n"
        "## Motivation\n(why this paper fits the team's direction)\n\n"
        "## Implementation plan\n(3-5 concrete steps a coding agent would follow)\n\n"
        "## Open questions\n(1-2 items needing clarification)"
    )
    messages = [{"role": "user", "content": user_prompt}]
    # Two-step tokenization: apply_chat_template(tokenize=False) returns a
    # string we can tokenize normally. In transformers 5.x,
    # apply_chat_template(tokenize=True, return_tensors="pt") returns a
    # BatchEncoding rather than a raw tensor, and BatchEncoding doesn't
    # expose .shape directly — generate() reads .shape on its first
    # positional arg and crashes. Going via string + tokenizer() yields
    # a BatchEncoding we can **unpack into generate().
    text = _tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = _tokenizer(text, return_tensors="pt").to(_model.device)
    prompt_len = inputs.input_ids.shape[1]

    with torch.no_grad():
        out = _model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            temperature=1.0,
            pad_token_id=_tokenizer.eos_token_id,
        )
    return _tokenizer.decode(out[0][prompt_len:], skip_special_tokens=True)