# ============================================================
# SPID — LLM extractor (OpenAI-compatible; default gpt-4o-mini)
#
# Uses the chat-completions JSON mode to extract structured
# project records.  Cost-controlled: only project-relevant
# windows of each section are sent to the model.
# ============================================================
from __future__ import annotations

import json
import re
import time

import httpx

from spid import taxonomy as tx
from spid.config import get_logger, get_settings
from spid.extract.base import ProjectMention
from spid.extract.rules import _PROJECT_GATE

log = get_logger("extract.llm")

_SYSTEM = """You are a financial analyst extracting CORPORATE PROJECTS from SEC filing text.
A "project" is any concrete initiative, investment, program, expansion, construction,
product/R&D effort, acquisition/integration, technology deployment, partnership, or
strategic program the company is pursuing.

Return STRICT JSON: {"projects": [ ... ]}. Each project object has:
  project_name  : short descriptive name (use the official name if one is stated)
  project_type  : EXACTLY ONE of the allowed types below
  description   : 1 sentence, grounded in the text (no speculation)
  objective     : the goal/why (or "")
  status        : one of planned | in_progress | completed | mentioned
  amount_usd    : numeric USD figure if stated (e.g. 1500000000), else null
  locations     : array of places (city/state/country), else []
  technologies  : array of technologies involved, else []
  partners      : array of named partners/JV counterparties, else []
  benefits      : array of expected benefits, else []
  risks         : array of stated risks, else []
  years         : array of 4-digit years mentioned, else []
  confidence    : 0.0-1.0 (how clearly this is a real, concrete project)

Allowed project_type values:
%s

Rules:
- Extract ONLY real initiatives, not generic risk-factor boilerplate or hypotheticals.
- If the text describes no concrete project, return {"projects": []}.
- Do not invent facts. Numbers/locations must appear in the text.
- Prefer fewer high-quality projects over many vague ones.""" % "\n".join(
    f"  - {k}: {v}" for k, v in tx.TYPE_LABELS.items()
)


def _focus(text: str, budget: int = 9000) -> str:
    """Keep only windows around project-signal hits to bound token cost."""
    if len(text) <= budget:
        return text
    spans: list[tuple[int, int]] = []
    for m in _PROJECT_GATE.finditer(text):
        spans.append((max(0, m.start() - 220), min(len(text), m.end() + 380)))
        if sum(b - a for a, b in spans) > budget * 1.4:
            break
    if not spans:
        return text[:budget]
    # merge overlapping windows
    spans.sort()
    merged = [spans[0]]
    for a, b in spans[1:]:
        if a <= merged[-1][1]:
            merged[-1] = (merged[-1][0], max(merged[-1][1], b))
        else:
            merged.append((a, b))
    chunks, total = [], 0
    for a, b in merged:
        chunks.append(text[a:b])
        total += b - a
        if total >= budget:
            break
    return "\n...\n".join(chunks)


class LLMExtractor:
    name = "llm"

    def __init__(self) -> None:
        self.s = get_settings()
        if not self.s.llm_api_key:
            raise RuntimeError("No LLM API key configured (OPENAI_API_KEY / MACLUSTR_API_KEY).")
        self._client = httpx.Client(
            base_url=self.s.llm_base_url,
            headers={"Authorization": f"Bearer {self.s.llm_api_key}"},
            timeout=self.s.llm_timeout,
        )

    def _call(self, content: str, retries: int = 3) -> dict:
        payload = {
            "model": self.s.llm_model,
            "messages": [
                {"role": "system", "content": _SYSTEM},
                {"role": "user", "content": content},
            ],
            "temperature": 0,
            "response_format": {"type": "json_object"},
            "max_tokens": 2000,
        }
        last = None
        for attempt in range(retries):
            try:
                r = self._client.post("/chat/completions", json=payload)
                if r.status_code == 429 or r.status_code >= 500:
                    raise httpx.HTTPStatusError("retryable", request=r.request, response=r)
                r.raise_for_status()
                txt = r.json()["choices"][0]["message"]["content"]
                return json.loads(txt)
            except (httpx.HTTPError, json.JSONDecodeError, KeyError) as e:
                last = e
                time.sleep(1.5 * (attempt + 1))
        log.warning("LLM call failed after %d retries: %s", retries, last)
        return {"projects": []}

    def extract(self, *, text: str, accession_number: str, cik: str,
                section_id: str, section_name: str) -> list[ProjectMention]:
        if not text or not _PROJECT_GATE.search(text):
            return []
        content = (f"SECTION: {section_name}\n\n" + _focus(text))[:14000]
        data = self._call(content)
        out: list[ProjectMention] = []
        for p in data.get("projects", []) or []:
            if not isinstance(p, dict):
                continue
            ptype = self._norm_type(p.get("project_type", ""))
            desc = str(p.get("description", "")).strip()[:600]
            if not desc:
                continue
            amount = p.get("amount_usd")
            try:
                amount = float(amount) if amount not in (None, "", "null") else None
            except (TypeError, ValueError):
                amount = None
            anchor = f"{ptype}:{str(p.get('project_name',''))[:40]}"
            conf = p.get("confidence", 0.6)
            try:
                conf = float(conf)
            except (TypeError, ValueError):
                conf = 0.6
            out.append(ProjectMention(
                mention_id=ProjectMention.make_id(accession_number, section_id, ptype, anchor),
                accession_number=accession_number, cik=cik,
                section_id=section_id, section_name=section_name,
                project_name=str(p.get("project_name") or tx.TYPE_LABELS.get(ptype, ptype))[:200],
                project_type=ptype, description=desc,
                objective=str(p.get("objective", ""))[:400],
                amount_usd=amount,
                amount_raw=(f"${amount:,.0f}" if amount else None),
                locations=_strlist(p.get("locations")),
                technologies=_strlist(p.get("technologies")) or tx.extract_technologies(desc),
                partners=_strlist(p.get("partners")),
                benefits=_strlist(p.get("benefits")),
                risks=_strlist(p.get("risks")),
                years=_intlist(p.get("years")) or tx.extract_years(desc),
                status=self._norm_status(p.get("status")),
                confidence=round(max(0.0, min(1.0, conf)), 3),
                backend=self.name,
            ))
        return out

    @staticmethod
    def _norm_type(raw: str) -> str:
        raw = (raw or "").strip().lower().replace(" ", "_").replace("-", "_")
        if raw in tx.PROJECT_TYPES:
            return raw
        # fuzzy contains match against labels/keys
        for k, label in tx.TYPE_LABELS.items():
            if raw and (raw in k or raw in label.lower().replace(" ", "_")):
                return k
        return "technology_deployment"

    @staticmethod
    def _norm_status(raw) -> str:
        raw = (str(raw or "")).strip().lower()
        return raw if raw in ("planned", "in_progress", "completed", "mentioned") else "mentioned"


def _strlist(v) -> list[str]:
    if not v:
        return []
    if isinstance(v, str):
        v = [v]
    return [str(x).strip()[:120] for x in v if str(x).strip()][:8]


def _intlist(v) -> list[int]:
    out = []
    for x in (v or []):
        try:
            y = int(str(x)[:4])
            if 1980 <= y <= 2040:
                out.append(y)
        except (TypeError, ValueError):
            continue
    return sorted(set(out))