from __future__ import annotations

import asyncio
import hashlib
import json
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_QUALIFIER
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


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()


DEFAULT_CONTRACT_PROMPT = """
OUTPUT SCHEMA (return array matching input order):
{"results":[{"id":123,"score":8,"decision":"accepted","reason":"до 10 слов на русском","reject_tag":null}]}

Rules:
- Input is a JSON array of posts.
- Return JSON only. No markdown. No text outside JSON.
- score must be integer 1..10.
- decision must be exactly one of: "accepted", "rejected", "maybe".
- reject_tag (rejected/maybe only) must be one of: "meme", "no_product", "politics", "discount_only", "off_topic", "vacancy", "low_content", "wrong_language", "injection", "weapon", null.
- Return one result for every input post id.
"""


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_results(data: dict, expected_ids: set[int], min_score: int) -> list[dict]:
    results = data.get("results")
    if not isinstance(results, list) and isinstance(data.get("data"), dict):
        results = data["data"].get("results")
    if not isinstance(results, list):
        raise ValueError("AI response has no results list")

    out: list[dict] = []
    seen: set[int] = set()
    for item in results:
        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}")
        score = max(1, min(10, int(item.get("score"))))
        decision = str(item.get("decision") or "").strip().lower()
        if decision not in {"accepted", "rejected", "maybe"}:
            decision = "accepted" if score >= min_score else "rejected"
        reason = str(item.get("reason") or "").strip()[:1000]
        reject_tag = item.get("reject_tag")
        reject_tag = str(reject_tag).strip().lower()[:80] if reject_tag not in {None, ""} else None
        seen.add(post_id)
        out.append({"id": post_id, "score": score, "decision": decision, "reason": reason, "reject_tag": reject_tag})

    missing = expected_ids - seen
    if missing:
        raise ValueError(f"AI response missing ids: {sorted(missing)[:10]}")
    return out


class AIQualifierWorker:
    def __init__(self) -> None:
        self.pool = None
        self.heartbeat = HeartbeatReporter(WORKER_AI_QUALIFIER, 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 COALESCE(qualification_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 COALESCE(qualification_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 created_at ASC, id ASC
                      FOR UPDATE SKIP LOCKED
                      LIMIT $1
                    )
                    UPDATE raw_posts rp
                    SET qualification_status='processing',
                        updated_at=NOW()
                    FROM cte
                    WHERE rp.id=cte.id
                    RETURNING rp.id, rp.raw_text, rp.original_url, rp.storage_post_url, rp.created_at,
                              rp.source_id,
                              (SELECT s.name FROM sources s WHERE s.id=rp.source_id) AS source_name,
                              (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_qualification_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),
                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:
        counts = {"accepted": 0, "rejected": 0, "maybe": 0}
        if isinstance(response, dict) and isinstance(response.get("results"), list):
            for item in response["results"]:
                decision = str(item.get("decision") or "")
                if decision in counts:
                    counts[decision] += 1
        usage = usage or {}
        await self.pool.execute(
            """
            UPDATE ai_qualification_batches
            SET status=$2,
                response_json=$3::jsonb,
                error=$4,
                accepted_count=$5,
                rejected_count=$6,
                maybe_count=$7,
                prompt_tokens=$8,
                completion_tokens=$9,
                total_tokens=$10,
                estimated_cost_usd=$11,
                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,
            counts["accepted"],
            counts["rejected"],
            counts["maybe"],
            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 qualification_status='failed',
                qualification_reason=$2,
                updated_at=NOW()
            WHERE id=ANY($1::bigint[])
            """,
            post_ids,
            error[:1000],
        )

    async def apply_results(self, batch_id: int, results: list[dict], model: str, prompt: str, min_score: int) -> None:
        for item in results:
            score = int(item["score"])
            decision = str(item["decision"])
            model_decision = decision
            if score >= min_score and decision == "maybe":
                decision = "accepted"
            if score < min_score and decision == "accepted":
                decision = "rejected"
            status = "accepted" if score >= min_score and decision == "accepted" else "rejected"
            await self.pool.execute(
                """
                UPDATE raw_posts
                SET qualification_status=$2,
                    qualification_score=$3,
                    qualification_decision=$4,
                    qualification_reason=$5,
                    qualification_model=$6,
                    qualification_prompt_hash=$7,
                    qualification_batch_id=$8,
                    qualification_reject_tag=$9,
                    qualification_model_decision=$10,
                    qualified_at=NOW(),
                    updated_at=NOW()
                WHERE id=$1
                """,
                int(item["id"]),
                status,
                score,
                decision,
                item["reason"],
                model,
                prompt_hash(prompt),
                batch_id,
                item.get("reject_tag"),
                model_decision,
            )

    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_QUALIFIER)
        setting_enabled = await fetch_bool_setting("ai_qualifier_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_qualifier_provider", "openrouter") or "openrouter")
        model = str(await fetch_setting("ai_qualifier_model", "") or "").strip()
        api_key = str(await fetch_setting("ai_qualifier_api_key", "") or "").strip()
        api_base = str(await fetch_setting("ai_qualifier_api_base", "") or "").strip()
        batch_size = max(1, await fetch_int_setting("ai_qualifier_batch_size", 30))
        min_score = max(1, min(10, await fetch_int_setting("ai_qualifier_min_score", 7)))
        max_text_chars = max(100, await fetch_int_setting("ai_qualifier_max_text_chars", 2000))
        temperature = max(0.0, await fetch_float_setting("ai_qualifier_temperature", 0.0))
        timeout = max(10, await fetch_int_setting("ai_qualifier_timeout_sec", 120))
        prompt = build_system_prompt(
            str(await fetch_setting("ai_qualifier_prompt", "") or ""),
            str(await fetch_setting("ai_qualifier_contract", DEFAULT_CONTRACT_PROMPT) or DEFAULT_CONTRACT_PROMPT),
        )

        if not model or not api_key or not prompt:
            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 = [
            {
                "id": int(p["id"]),
                "source": p.get("source_name") or "",
                "original_url": p.get("original_url") 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]
        batch_id = await self.create_batch(provider, model, prompt, payload)
        try:
            response, usage = await self.call_ai(provider, model, api_key, api_base, prompt, payload, temperature, timeout)
            results = validate_results(response, set(post_ids), min_score)
            await self.apply_results(batch_id, results, normalize_model(provider, model), prompt, min_score)
            await self.finish_batch(batch_id, "done", {"results": results}, usage=usage)
            logger.info("AI qualifier batch done: id={} posts={} usage={}", batch_id, len(results), usage)
            return True
        except Exception as exc:
            await self.mark_posts_failed(post_ids, str(exc))
            await self.finish_batch(batch_id, "failed", error=str(exc))
            logger.exception("AI qualifier batch failed: id={} error={}", batch_id, exc)
            return True

    async def run_loop(self) -> None:
        await self.init()
        worker_id = f"{WORKER_AI_QUALIFIER}"
        logger.info("{} started", worker_id)
        while True:
            try:
                had_work = await self.run_once()
            except Exception as exc:
                logger.exception("AI qualifier loop error: {}", exc)
                had_work = False
            interval = max(10, await fetch_int_setting("ai_qualifier_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 = AIQualifierWorker()
    await worker.run_loop()


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