from __future__ import annotations

import asyncio
import hashlib
import json
import re
from datetime import datetime, timezone
from decimal import Decimal
from typing import Any

import litellm
from loguru import logger

from ..config import settings
from ..constants import WORKER_AI_WRITER
from ..db import fetch_bool_setting, fetch_float_setting, fetch_int_setting, fetch_setting, get_pool
from ..heartbeat import HeartbeatReporter
from ..jobs import is_worker_enabled
from ..text_utils import build_publication_text, normalize_hash_tag, parse_categories


NON_TARGET_CATEGORY_ID = 18
NON_TARGET_CATEGORY_TAGS = {"нцк", "nck"}
EMOJI_LEADING_MARKER_RE = re.compile(
    r"^[ \t]*(?:🧵|🎯|🛡|🧤|📦|💵|⚙\ufe0f?|🎨|🎒|🔧|👕|💪)\s+"
)


def now_utc() -> datetime:
    return datetime.now(timezone.utc)


def prompt_hash(prompt: str) -> str:
    return hashlib.sha256((prompt or "").encode("utf-8")).hexdigest()


def normalize_prompt(prompt: str) -> str:
    return (prompt or "").replace("\\r\\n", "\n").replace("\\n", "\n").strip()


def strip_leading_emoji_markers(text: str) -> str:
    lines = str(text or "").splitlines()
    cleaned: list[str] = []
    for line in lines:
        cleaned.append(EMOJI_LEADING_MARKER_RE.sub("", line))
    return "\n".join(cleaned).strip()


DEFAULT_CONTRACT_PROMPT = """
OUTPUT SCHEMA (return array matching input order):
{"rewrites":[{"id":123,"category_id":2,"text":"готовый текст без хэштегов","notes":"короткая заметка для редактора или null"}]}

Rules:
- Input is a JSON object with key "posts" containing accepted posts.
- Input has key "categories" with objects: id, name, tag.
- Each post includes producer_name and producer_tag. producer_name can be a manufacturer, shop, or publishing source.
- The "text" field for every rewrite must explicitly mention producer_name at least once. Do not satisfy this only through producer_tag, hashtags, notes, or metadata.
- Pick exactly one category from the categories list in input and return its numeric id as category_id.
- If the post is non-target content, choose category_id 18 ("Не целевой контент", tag "НЦК"); the app will reject it automatically.
- Do not add links or hashtags to rewritten text.
- Return JSON only. No markdown. No text outside JSON.
- Never return an empty object. Never omit the "rewrites" key.
- Return one rewrite for every input post id.
- The "text" field must be a ready-to-publish Russian post with normal paragraph line breaks encoded as JSON string newlines.
- The "notes" field must be a short editor note in Russian or null.
"""


def build_system_prompt(user_prompt: str, contract_prompt: str) -> str:
    parts = [normalize_prompt(user_prompt), normalize_prompt(contract_prompt or DEFAULT_CONTRACT_PROMPT)]
    return "\n\n".join(part for part in parts if part).strip()


def normalize_model(provider: str, model: str) -> str:
    provider = (provider or "").strip().lower()
    model = (model or "").strip()
    if not model:
        return model
    if "/" in model:
        return model
    if provider in {"openrouter", "anthropic", "gemini", "vertex_ai", "bedrock"}:
        return f"{provider}/{model}"
    return model


def parse_ai_json(content: str) -> dict:
    raw = (content or "").strip()
    if raw.startswith("```"):
        raw = raw.strip("`")
        if raw.lower().startswith("json"):
            raw = raw[4:].strip()
    return json.loads(raw)


def response_usage(response: Any) -> dict[str, Any]:
    usage = getattr(response, "usage", None)
    if usage is None and isinstance(response, dict):
        usage = response.get("usage")

    def get(name: str) -> int | None:
        if usage is None:
            return None
        value = getattr(usage, name, None)
        if value is None and isinstance(usage, dict):
            value = usage.get(name)
        try:
            return int(value) if value is not None else None
        except Exception:
            return None

    cost = getattr(response, "_hidden_params", None)
    if isinstance(cost, dict):
        cost = cost.get("response_cost")
    else:
        cost = None
    try:
        cost_value = Decimal(str(cost)) if cost is not None else None
    except Exception:
        cost_value = None

    return {
        "prompt_tokens": get("prompt_tokens"),
        "completion_tokens": get("completion_tokens"),
        "total_tokens": get("total_tokens"),
        "estimated_cost_usd": cost_value,
    }


def validate_rewrites(
    data: dict,
    expected_ids: set[int],
    categories: list[dict[str, Any]],
    source_tags_by_id: dict[int, str],
) -> list[dict]:
    rewrites = data.get("rewrites")
    if not isinstance(rewrites, list) and isinstance(data.get("data"), dict):
        rewrites = data["data"].get("rewrites")
    if not isinstance(rewrites, list):
        raise ValueError("AI response has no rewrites list")
    out: list[dict] = []
    seen: set[int] = set()
    for item in rewrites:
        if not isinstance(item, dict):
            continue
        post_id = int(item.get("id"))
        if post_id not in expected_ids:
            raise ValueError(f"AI returned unexpected id={post_id}")
        text = strip_leading_emoji_markers(str(item.get("text") or "").strip())
        if len(text) < 40:
            raise ValueError(f"AI returned too short rewrite for id={post_id}")
        by_id = {int(c["id"]): c for c in categories if c.get("id") is not None}
        by_name = {str(c["name"]).lower(): c for c in categories if c.get("name")}
        if item.get("category_id") is not None:
            category_id = int(item["category_id"])
            if category_id not in by_id:
                raise ValueError(f"AI returned unknown category_id={category_id!r} for id={post_id}")
            category_row = by_id[category_id]
        else:
            category_name = str(item.get("category") or "").strip().strip("#")
            category_row = by_name.get(category_name.lower())
            if not category_row:
                raise ValueError(f"AI returned unknown category={category_name!r} for id={post_id}")
        notes = str(item.get("notes") or "").strip()[:1000]
        seen.add(post_id)
        category_tag = normalize_hash_tag(category_row.get("tag") or category_row.get("name") or "", "category")
        reject_by_category = int(category_row["id"]) == NON_TARGET_CATEGORY_ID or category_tag.lower() in NON_TARGET_CATEGORY_TAGS
        out.append(
            {
                "id": post_id,
                "category_id": int(category_row["id"]),
                "category": str(category_row["name"]),
                "category_tag": category_tag,
                "source_tag": normalize_hash_tag(source_tags_by_id.get(post_id) or "", "source"),
                "text": text,
                "notes": notes,
                "reject_by_category": reject_by_category,
            }
        )
    missing = expected_ids - seen
    if missing:
        raise ValueError(f"AI response missing ids: {sorted(missing)[:10]}")
    return out


async def load_writer_categories(pool) -> list[dict[str, Any]]:
    try:
        rows = await pool.fetch(
            """
            SELECT sort_order AS id, name, tag
            FROM content_categories
            WHERE is_active=TRUE
            ORDER BY sort_order, name
            """
        )
    except Exception:
        rows = []
    categories = [{"id": int(row["id"]), "name": str(row["name"]), "tag": str(row["tag"])} for row in rows]
    if categories:
        return categories
    legacy = parse_categories(await fetch_setting("ai_writer_categories", []))
    names = legacy or [
        "защита",
        "одежда",
        "разгрузка",
        "рюкзаки",
        "airsoft",
        "патчи",
        "электроника",
        "аксессуары",
        "производство",
    ]
    return [{"id": idx + 1, "name": name, "tag": normalize_hash_tag(name, "category")} for idx, name in enumerate(names)]


class AIWriterWorker:
    def __init__(self) -> None:
        self.pool = None
        self.heartbeat = HeartbeatReporter(WORKER_AI_WRITER, 30)

    async def init(self) -> None:
        self.pool = await get_pool()

    async def claim_posts(self, batch_size: int) -> list[dict]:
        async with self.pool.acquire() as conn:
            async with conn.transaction():
                count = await conn.fetchval(
                    """
                    SELECT COUNT(*)
                    FROM raw_posts
                    WHERE status='storage_ready'
                      AND qualification_status='accepted'
                      AND COALESCE(rewrite_status, 'pending') IN ('pending', 'failed')
                      AND EXISTS (
                        SELECT 1
                        FROM raw_post_media m
                        WHERE m.raw_post_id=raw_posts.id
                          AND COALESCE(m.editor_hidden, FALSE)=FALSE
                          AND m.media_type IN ('photo', 'video')
                          AND COALESCE(m.tg_file_id, m.storage_attachment_id, '') <> ''
                      )
                      AND NOT EXISTS (
                        SELECT 1
                        FROM raw_post_media m
                        WHERE m.raw_post_id=raw_posts.id
                          AND COALESCE(m.editor_hidden, FALSE)=FALSE
                          AND m.media_type IN ('photo', 'video')
                          AND COALESCE(m.tg_file_id, m.storage_attachment_id, '') = ''
                      )
                    """
                )
                if int(count or 0) < batch_size:
                    return []

                rows = await conn.fetch(
                    """
                    WITH cte AS (
                      SELECT id
                      FROM raw_posts
                      WHERE status='storage_ready'
                        AND qualification_status='accepted'
                        AND COALESCE(rewrite_status, 'pending') IN ('pending', 'failed')
                        AND EXISTS (
                          SELECT 1
                          FROM raw_post_media m
                          WHERE m.raw_post_id=raw_posts.id
                            AND COALESCE(m.editor_hidden, FALSE)=FALSE
                            AND m.media_type IN ('photo', 'video')
                            AND COALESCE(m.tg_file_id, m.storage_attachment_id, '') <> ''
                        )
                        AND NOT EXISTS (
                          SELECT 1
                          FROM raw_post_media m
                          WHERE m.raw_post_id=raw_posts.id
                            AND COALESCE(m.editor_hidden, FALSE)=FALSE
                            AND m.media_type IN ('photo', 'video')
                            AND COALESCE(m.tg_file_id, m.storage_attachment_id, '') = ''
                        )
                      ORDER BY qualified_at ASC NULLS LAST, id ASC
                      FOR UPDATE SKIP LOCKED
                      LIMIT $1
                    )
                    UPDATE raw_posts rp
                    SET rewrite_status='processing',
                        updated_at=NOW()
                    FROM cte
                    WHERE rp.id=cte.id
                    RETURNING rp.id, rp.raw_text, rp.original_url, rp.created_at, rp.qualification_score,
                              rp.qualification_reason,
                              (SELECT s.name FROM sources s WHERE s.id=rp.source_id) AS source_name,
                              (SELECT s.tag FROM sources s WHERE s.id=rp.source_id) AS source_tag,
                              (SELECT COUNT(*) FROM raw_post_media m WHERE m.raw_post_id=rp.id) AS media_count,
                              (SELECT ARRAY_AGG(DISTINCT m.media_type ORDER BY m.media_type)
                               FROM raw_post_media m WHERE m.raw_post_id=rp.id) AS media_types
                    """,
                    batch_size,
                )
                return [dict(r) for r in rows]

    async def create_batch(self, provider: str, model: str, prompt: str, payload: dict) -> int:
        return int(
            await self.pool.fetchval(
                """
                INSERT INTO ai_writer_batches(provider, model, prompt_hash, prompt_text, posts_count, request_json)
                VALUES($1, $2, $3, $4, $5, $6::jsonb)
                RETURNING id
                """,
                provider,
                model,
                prompt_hash(prompt),
                prompt,
                len(payload["posts"]),
                json.dumps(payload, ensure_ascii=False),
            )
        )

    async def finish_batch(
        self,
        batch_id: int,
        status: str,
        response: Any = None,
        error: str | None = None,
        usage: dict[str, Any] | None = None,
    ) -> None:
        ready_count = 0
        if isinstance(response, dict) and isinstance(response.get("rewrites"), list):
            ready_count = len(response["rewrites"])
        usage = usage or {}
        await self.pool.execute(
            """
            UPDATE ai_writer_batches
            SET status=$2,
                response_json=$3::jsonb,
                error=$4,
                ready_count=$5,
                failed_count=$6,
                prompt_tokens=$7,
                completion_tokens=$8,
                total_tokens=$9,
                estimated_cost_usd=$10,
                completed_at=NOW()
            WHERE id=$1
            """,
            batch_id,
            status,
            json.dumps(response, ensure_ascii=False) if response is not None else None,
            error[:2000] if error else None,
            ready_count if status == "done" else 0,
            0 if status == "done" else 1,
            usage.get("prompt_tokens"),
            usage.get("completion_tokens"),
            usage.get("total_tokens"),
            usage.get("estimated_cost_usd"),
        )

    async def mark_posts_failed(self, post_ids: list[int], error: str) -> None:
        if not post_ids:
            return
        await self.pool.execute(
            """
            UPDATE raw_posts
            SET rewrite_status='failed',
                rewrite_notes=$2,
                updated_at=NOW()
            WHERE id=ANY($1::bigint[])
            """,
            post_ids,
            error[:1000],
        )

    async def apply_rewrites(self, batch_id: int, rewrites: list[dict], model: str, prompt: str) -> None:
        for item in rewrites:
            final_text = build_publication_text(item["text"], item["category_tag"], item["source_tag"])
            editorial_status = "rejected" if item.get("reject_by_category") else "review"
            editor_notes = "Отклонено AI: Не целевой контент" if item.get("reject_by_category") else None
            await self.pool.execute(
                """
                UPDATE raw_posts
                SET rewrite_status='ready',
                    rewritten_text=$2,
                    rewrite_notes=$3,
                    rewrite_model=$4,
                    rewrite_prompt_hash=$5,
                    rewrite_batch_id=$6,
                    rewrite_category=$7,
                    rewrite_category_tag=$8,
                    rewrite_source_tag=$9,
                    rewrite_category_id=$10,
                    editorial_status=$11,
                    final_text=$12,
                    final_category=$7,
                    final_category_tag=$8,
                    final_source_tag=$9,
                    final_category_id=$10,
                    editor_notes=$13,
                    qualification_status=CASE WHEN $11='rejected' THEN 'rejected' ELSE qualification_status END,
                    qualification_decision=CASE WHEN $11='rejected' THEN 'rejected' ELSE qualification_decision END,
                    qualification_model_decision=CASE WHEN $11='rejected' THEN 'rejected' ELSE qualification_model_decision END,
                    qualification_reason=CASE WHEN $11='rejected' THEN 'AI writer selected non-target category' ELSE qualification_reason END,
                    qualification_reject_tag=CASE WHEN $11='rejected' THEN 'non_target_content' ELSE qualification_reject_tag END,
                    reviewed_at=CASE WHEN $11='rejected' THEN NOW() ELSE reviewed_at END,
                    rewritten_at=NOW(),
                    updated_at=NOW()
                WHERE id=$1
                """,
                int(item["id"]),
                item["text"],
                item["notes"],
                model,
                prompt_hash(prompt),
                batch_id,
                item["category"],
                item["category_tag"],
                item["source_tag"],
                item["category_id"],
                editorial_status,
                final_text,
                editor_notes,
            )

    async def call_ai(
        self,
        provider: str,
        model: str,
        api_key: str,
        api_base: str,
        prompt: str,
        payload: dict,
        temperature: float,
        timeout: int,
    ) -> tuple[dict, dict[str, Any]]:
        messages = [
            {"role": "system", "content": prompt},
            {"role": "user", "content": json.dumps(payload, ensure_ascii=False)},
        ]
        kwargs: dict[str, Any] = {
            "model": normalize_model(provider, model),
            "messages": messages,
            "temperature": temperature,
            "timeout": timeout,
            "response_format": {"type": "json_object"},
        }
        if api_key:
            kwargs["api_key"] = api_key
        if api_base:
            kwargs["api_base"] = api_base
        response = await asyncio.to_thread(litellm.completion, **kwargs)
        content = response.choices[0].message.content
        return parse_ai_json(content), response_usage(response)

    async def run_once(self) -> bool:
        enabled = await is_worker_enabled(self.pool, WORKER_AI_WRITER)
        setting_enabled = await fetch_bool_setting("ai_writer_enabled", False)
        if not enabled or not setting_enabled:
            await self.heartbeat.beat(self.pool, status="disabled", force=True)
            return False

        provider = str(await fetch_setting("ai_writer_provider", "anthropic") or "anthropic")
        model = str(await fetch_setting("ai_writer_model", "") or "").strip()
        api_key = str(await fetch_setting("ai_writer_api_key", "") or "").strip()
        api_base = str(await fetch_setting("ai_writer_api_base", "") or "").strip()
        batch_size = max(1, await fetch_int_setting("ai_writer_batch_size", 1))
        max_text_chars = max(200, await fetch_int_setting("ai_writer_max_text_chars", 3500))
        temperature = max(0.0, await fetch_float_setting("ai_writer_temperature", 0.4))
        timeout = max(10, await fetch_int_setting("ai_writer_timeout_sec", 180))
        prompt = build_system_prompt(
            str(await fetch_setting("ai_writer_prompt", "") or ""),
            str(await fetch_setting("ai_writer_contract", DEFAULT_CONTRACT_PROMPT) or DEFAULT_CONTRACT_PROMPT),
        )
        categories = await load_writer_categories(self.pool)

        if not model or not api_key or not prompt or not categories:
            await self.heartbeat.beat(self.pool, status="not_configured", force=True)
            return False

        posts = await self.claim_posts(batch_size)
        await self.heartbeat.beat(self.pool, meta={"claimed": len(posts), "batch_size": batch_size})
        if not posts:
            return False

        payload = {
            "task": "rewrite_accepted_posts",
            "categories": categories,
            "posts": [
                {
                    "id": int(p["id"]),
                    "producer_name": p.get("source_name") or "",
                    "producer_tag": normalize_hash_tag(p.get("source_tag") or p.get("source_name") or "", "source"),
                    "original_url": p.get("original_url") or "",
                    "qualification_score": p.get("qualification_score"),
                    "qualification_reason": p.get("qualification_reason") or "",
                    "media_count": int(p.get("media_count") or 0),
                    "media_types": list(p.get("media_types") or []),
                    "text": str(p.get("raw_text") or "")[:max_text_chars],
                }
                for p in posts
            ],
        }
        post_ids = [int(p["id"]) for p in posts]
        source_tags_by_id = {
            int(p["id"]): normalize_hash_tag(p.get("source_tag") or p.get("source_name") or "", "source")
            for p in posts
        }
        batch_id = await self.create_batch(provider, model, prompt, payload)
        response = None
        usage = None
        try:
            response, usage = await self.call_ai(provider, model, api_key, api_base, prompt, payload, temperature, timeout)
            rewrites = validate_rewrites(response, set(post_ids), categories, source_tags_by_id)
            await self.apply_rewrites(batch_id, rewrites, normalize_model(provider, model), prompt)
            await self.finish_batch(batch_id, "done", {"rewrites": rewrites}, usage=usage)
            logger.info("AI writer batch done: id={} posts={} usage={}", batch_id, len(rewrites), usage)
            return True
        except Exception as exc:
            await self.mark_posts_failed(post_ids, str(exc))
            await self.finish_batch(batch_id, "failed", response=response, error=str(exc), usage=usage)
            logger.exception("AI writer batch failed: id={} error={}", batch_id, exc)
            return True

    async def run_loop(self) -> None:
        await self.init()
        logger.info("{} started", WORKER_AI_WRITER)
        while True:
            try:
                had_work = await self.run_once()
            except Exception as exc:
                logger.exception("AI writer loop error: {}", exc)
                had_work = False
            interval = max(10, await fetch_int_setting("ai_writer_interval_sec", 60))
            if not had_work:
                await asyncio.sleep(interval)


async def main() -> None:
    logger.remove()
    logger.add(lambda msg: print(msg, end=""), level=settings.log_level)
    worker = AIWriterWorker()
    await worker.run_loop()


if __name__ == "__main__":
    asyncio.run(main())
