"""Classify service — LLM-based document classification."""
import json, re, os
from concurrent.futures import ThreadPoolExecutor, as_completed

import httpx
from db import documents as db_docs
from llm_prompt import build_classify_prompt

log = __import__("logging").getLogger(__name__)

MAX_WORKERS = 4
LLM_URL = "https://api.aillm.ru/v1/chat/completions"
LLM_KEY = os.environ.get("LLM_KEY") or os.environ.get("LLM_API_KEY", "")
LLM_MODEL = "gpt-oss-120b"


def _call_llm_classify(header_text):
    """Call LLM for classification. Returns parsed JSON dict."""
    prompt, _ = build_classify_prompt(header_text)
    payload = {
        "model": LLM_MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 500,
        "temperature": 0.1,
    }
    with httpx.Client(http2=True, timeout=60, verify=False) as client:
        resp = client.post(
            LLM_URL, json=payload,
            headers={"Authorization": f"Bearer {LLM_KEY}", "Content-Type": "application/json"},
        )
        resp.raise_for_status()
        data = resp.json()

    raw = data.get("choices", [{}])[0].get("message", {}).get("content", "")
    # Extract JSON from possible markdown wrapping
    json_text = raw
    if "```json" in json_text:
        json_text = json_text.split("```json")[1].split("```")[0]
    elif "```" in json_text:
        json_text = json_text.split("```")[1].split("```")[0]
    return json.loads(json_text.strip())


def classify_batch(batch_id):
    """Classify all pending documents in a batch. Returns summary dict."""
    pending = db_docs.list_pending(batch_id)
    if not pending:
        return {"ok": False, "error": "no pending documents"}

    total = len(pending)
    done = 0
    failed = 0

    def _classify_one(doc):
        try:
            text = _smart_extract(doc["elements_json"])
            result = _call_llm_classify(text)
            db_docs.set_classification(
                doc["id"],
                result.get("doc_type", "other"),
                result.get("own_number"),
                result.get("parent_number"),
                result.get("doc_date"),
                result.get("counterparty"),
            )
            return True
        except Exception as e:
            db_docs.set_classify_failed(doc["id"], str(e))
            return False

    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
        futures = {pool.submit(_classify_one, d): d for d in pending}
        for f in as_completed(futures):
            if f.result():
                done += 1
            else:
                failed += 1

    return {"ok": True, "total": total, "done": done, "failed": failed}


def _smart_extract(elements_json):
    """Extract smart header: first ~1500 chars + marker lines from full doc."""
    if not elements_json:
        return ""

    if isinstance(elements_json, str):
        try:
            elements = json.loads(elements_json)
        except json.JSONDecodeError:
            return elements_json[:2000]
    elif isinstance(elements_json, list):
        elements = elements_json
    else:
        return str(elements_json)[:2000]

    # Build full text
    lines = []
    for el in elements:
        if isinstance(el, dict):
            t = el.get("type") or el.get("TYPE", "")
            if t == "paragraph":
                txt = el.get("text") or el.get("TEXT", "")
                if txt:
                    lines.append(txt)
            elif t == "table":
                rows = el.get("rows") or el.get("ROWS", [])
                for row in rows:
                    lines.append(" | ".join(str(c) for c in row))

    full_text = "\n".join(lines)

    # Header: first ~1500 chars
    header = full_text[:1500]

    # Marker lines: grep for key patterns
    markers = re.findall(
        r'.{0,200}(?:договор|№|соглашен|приложен|специф|контрагент|заказчик|арендатор).{0,200}',
        full_text, re.IGNORECASE,
    )
    unique_markers = list(dict.fromkeys(markers))[:10]

    combined = header + "\n---\n" + "\n".join(unique_markers)
    return combined[:3000]