"""
ResearchOwl Processor
Chunking → Quality scoring via Ollama → Embeddings → RAG synthesis
"""
import asyncio
import json
import math
import re
from typing import Optional

import httpx
import structlog

from src.config import settings
from src.db.database import ResearchDB

logger = structlog.get_logger()


class OllamaClient:
    """Async client for Ollama API"""

    def __init__(self):
        self.base_url = settings.ollama_url.rstrip("/")
        self.model = settings.ollama_model
        self.embed_model = settings.ollama_embed_model

    async def generate(self, prompt: str, system: str = None,
                       timeout: int = 120, temperature: float = 0.7) -> str:
        payload = {
            "model": self.model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": temperature,
                "num_predict": 2048,
                "repeat_penalty": 1.15,
                "repeat_last_n": 128,
            }
        }
        if system:
            payload["system"] = system

        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.post(f"{self.base_url}/api/generate", json=payload)
            resp.raise_for_status()
            return resp.json().get("response", "").strip()

    async def embed(self, text: str) -> Optional[list[float]]:
        """Get embedding vector for a text"""
        payload = {"model": self.embed_model, "prompt": text}
        try:
            async with httpx.AsyncClient(timeout=60) as client:
                resp = await client.post(f"{self.base_url}/api/embeddings", json=payload)
                resp.raise_for_status()
                return resp.json().get("embedding")
        except Exception as e:
            logger.warning("Embedding failed", error=str(e))
            return None

    async def is_available(self) -> bool:
        try:
            async with httpx.AsyncClient(timeout=5) as client:
                resp = await client.get(f"{self.base_url}/api/tags")
                return resp.status_code == 200
        except Exception:
            return False


def simple_chunk(text: str, chunk_size: int = 800, overlap: int = 100) -> list[str]:
    """
    Split text into overlapping chunks by approximate word count.
    Respects paragraph/line boundaries when possible.

    Acepta párrafos separados por uno o más saltos de línea (Wikipedia y
    trafilatura usan '\n' simple, lo que antes dejaba el documento entero como
    un único 'párrafo' → un solo chunk gigante). Además subdivide por palabras
    cualquier párrafo que por sí solo supere chunk_size.
    """
    raw_paragraphs = [p.strip() for p in re.split(r"\n+", text) if p.strip()]

    # Subdivide párrafos sobredimensionados en piezas de (chunk_size - overlap)
    # palabras; así, al reinyectar 'overlap' palabras de solapamiento, ningún
    # chunk resultante supera chunk_size.
    piece_size = max(1, chunk_size - max(0, overlap))
    paragraphs: list[str] = []
    for para in raw_paragraphs:
        words = para.split()
        if len(words) <= chunk_size:
            paragraphs.append(para)
        else:
            for i in range(0, len(words), piece_size):
                paragraphs.append(" ".join(words[i:i + piece_size]))

    chunks = []
    current = []
    current_words = 0

    for para in paragraphs:
        para_words = len(para.split())
        if current_words + para_words > chunk_size and current:
            chunks.append("\n\n".join(current))
            # overlap: arrastra un tail de 'overlap' palabras (no el párrafo
            # completo — eso duplicaba el tamaño cuando los párrafos eran grandes)
            if overlap > 0:
                tail = "\n\n".join(current).split()[-overlap:]
                current = [" ".join(tail)]
                current_words = len(tail)
            else:
                current = []
                current_words = 0
        current.append(para)
        current_words += para_words

    if current:
        chunks.append("\n\n".join(current))

    return chunks


def cosine_similarity(a: list[float], b: list[float]) -> float:
    """Simple cosine similarity"""
    if not a or not b or len(a) != len(b):
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)


class ContentProcessor:
    """
    Processes scraped sources:
    1. Chunks content
    2. Scores quality with Ollama
    3. Generates embeddings
    4. Stores high-quality chunks
    """

    def __init__(self, db: ResearchDB, ollama: OllamaClient):
        self.db = db
        self.ollama = ollama

    async def process_session(self, session_id: int, topic: str,
                               progress_callback=None) -> dict:
        """Process all scraped sources for a session"""
        sources = await self.db.get_all_sources(session_id)
        scraped = [s for s in sources if s["status"] == "scraped"]

        logger.info("Processing sources", total=len(scraped))
        scraped = await self._dedup_sources(session_id, scraped)
        logger.info("After dedup", unique=len(scraped))
        total_chunks = 0

        semaphore = asyncio.Semaphore(3)  # process 3 sources at once

        async def process_one(source):
            async with semaphore:
                n = await self._process_source(session_id, topic, source)
                return n

        results = await asyncio.gather(*[process_one(s) for s in scraped],
                                        return_exceptions=True)

        for i, r in enumerate(results):
            if isinstance(r, Exception):
                logger.error("Source processing raised exception",
                             source_id=scraped[i]["id"], error=str(r), exc_info=r)
            elif isinstance(r, int):
                total_chunks += r

        total_words = sum(s.get("word_count", 0) for s in scraped)
        await self.db.update_session(
            session_id,
            total_chunks=total_chunks,
            total_words=total_words
        )

        if progress_callback:
            await progress_callback(total_chunks=total_chunks, total_words=total_words)

        return {"total_chunks": total_chunks, "total_words": total_words}

    async def _dedup_sources(self, session_id: int,
                              scraped: list[dict]) -> list[dict]:
        try:
            import hashlib
            seen_hashes: set = set()
            seen_prefixes: list = []
            unique: list = []
            duplicates = 0

            for source in scraped:
                content = await self.db.get_source_content(source["id"])
                if not content:
                    unique.append(source)
                    continue

                content_hash = hashlib.md5(content[:2000].encode()).hexdigest()
                if content_hash in seen_hashes:
                    duplicates += 1
                    await self.db.update_source(source["id"], status="skipped")
                    continue
                seen_hashes.add(content_hash)

                prefix = content[:300].strip().lower()
                prefix_words = set(prefix.split())
                is_dup = False
                if len(prefix_words) >= 10:
                    for seen_prefix_words in seen_prefixes:
                        intersection = len(prefix_words & seen_prefix_words)
                        union = len(prefix_words | seen_prefix_words)
                        if intersection / max(union, 1) > 0.85:
                            is_dup = True
                            break

                if is_dup:
                    duplicates += 1
                    await self.db.update_source(source["id"], status="skipped")
                    continue

                seen_prefixes.append(prefix_words)
                unique.append(source)

            if duplicates > 0:
                logger.info("Dedup complete", session_id=session_id,
                            original=len(scraped), duplicates=duplicates,
                            unique=len(unique))
            return unique
        except Exception as e:
            logger.warning("Dedup failed, processing all sources", error=str(e))
            return scraped

    async def _process_source(self, session_id: int, topic: str, source: dict) -> int:
        """Chunk, score, embed and store a single source. Returns chunk count."""
        source_id = source["id"]

        content = await self.db.get_source_content(source_id)
        if not content:
            logger.warning("No content in source_contents", source_id=source_id)
            return 0

        chunks = simple_chunk(content, settings.chunk_size, settings.chunk_overlap)
        # Solo se puntúan/almacenan chunks con suficiente contenido
        candidates = [(i, ch) for i, ch in enumerate(chunks)
                      if len(ch.split()) >= 30]
        logger.info("Processing source", source_id=source_id,
                    content_len=len(content), num_chunks=len(chunks),
                    candidates=len(candidates),
                    quality_threshold=settings.quality_threshold)
        if not candidates:
            return 0

        # Scoring en lote: 1 (o pocas) llamadas a Claude por fuente en vez de
        # una por chunk — antes una fuente de 19 chunks = 19 llamadas.
        qualities = await self._score_quality_batch(
            [ch for _, ch in candidates], topic, session_id
        )

        stored = 0
        filtered_quality = 0

        for (i, chunk), quality in zip(candidates, qualities):
            if quality < settings.quality_threshold:
                filtered_quality += 1
                logger.debug("Chunk filtered by quality", source_id=source_id,
                             chunk_index=i, quality=round(quality, 2),
                             threshold=settings.quality_threshold,
                             words=len(chunk.split()))
                continue

            # Embeber el chunk completo (ya acotado a ~chunk_size palabras).
            # Antes truncaba a 1000 chars → el vector solo representaba el
            # principio de cada chunk, degradando el ranking del RAG.
            embedding = await self.ollama.embed(chunk)

            await self.db.add_chunk(
                session_id=session_id,
                source_id=source_id,
                content=chunk,
                chunk_index=i,
                token_count=len(chunk.split()),
                quality_score=quality,
                embedding=embedding
            )
            stored += 1

        if filtered_quality > 0 and stored == 0:
            logger.warning(
                "All chunks filtered by quality — consider lowering QUALITY_THRESHOLD "
                "(currently %.1f) or set QUALITY_THRESHOLD=0 to disable",
                settings.quality_threshold,
                source_id=source_id, chunks_total=len(chunks),
                chunks_filtered=filtered_quality
            )

        logger.info("Source processed", source_id=source_id, stored=stored)
        return stored

    async def _score_quality(self, chunk: str, topic: str,
                              session_id: int | None = None) -> float:
        """Score 0-1 relevance to topic. Uses Claude Haiku if API key set, else Ollama."""
        if settings.anthropic_api_key:
            return await self._score_with_claude(chunk, topic, session_id)
        return await self._score_with_ollama(chunk, topic)

    # ─── Batch scoring ──────────────────────────────────────────────────────

    _BATCH_SCORE_SIZE = 25

    async def _score_quality_batch(self, chunk_texts: list[str], topic: str,
                                    session_id: int | None = None) -> list[float]:
        """Puntúa varios chunks a la vez. Devuelve scores 0-1 en el mismo orden."""
        if not chunk_texts:
            return []
        if not settings.anthropic_api_key:
            # Ollama es local; no merece la pena batchear, se mantiene por-chunk
            return [await self._score_with_ollama(c, topic) for c in chunk_texts]

        results: list[float] = []
        for i in range(0, len(chunk_texts), self._BATCH_SCORE_SIZE):
            sub = chunk_texts[i:i + self._BATCH_SCORE_SIZE]
            scores = await self._score_with_claude_batch(sub, topic, session_id)
            if scores is None:  # fallo del batch → fallback por-chunk con Ollama
                scores = [await self._score_with_ollama(c, topic) for c in sub]
            results.extend(scores)
        return results

    @staticmethod
    def _parse_batch_scores(text: str, n: int) -> list[float]:
        """Extrae n scores normalizados (0-1) de la respuesta del modelo.

        Toma el último número de cada línea (robusto ante numeración tipo
        '1. 8'); rellena con 0.6 neutro si faltan líneas."""
        scores: list[float] = []
        for line in text.splitlines():
            nums = re.findall(r'\d+(?:\.\d+)?', line)
            if nums:
                scores.append(min(1.0, float(nums[-1]) / 10.0))
        if len(scores) >= n:
            return scores[:n]
        return scores + [0.6] * (n - len(scores))

    async def _score_with_claude_batch(self, chunks: list[str], topic: str,
                                        session_id: int | None = None):
        """Puntúa hasta _BATCH_SCORE_SIZE chunks en una sola llamada.
        Devuelve lista de scores 0-1, o None si la llamada falla."""
        from src.llm import get_anthropic_client
        listing = "\n\n".join(f"[{i + 1}]\n{c[:400]}" for i, c in enumerate(chunks))
        prompt = (
            f'Rate each of the following {len(chunks)} texts 0-10 for relevance '
            f'to the topic "{topic}". Be generous — tangentially related = 4+, '
            f'only below 3 if completely unrelated.\n'
            f'Reply with EXACTLY {len(chunks)} lines, one integer 0-10 per line, '
            f'in the same order as the texts. Just the number (e.g. 7), '
            f'no labels, no other text.\n\n{listing}'
        )
        try:
            client = get_anthropic_client()
            msg = await client.messages.create(
                model=settings.claude_model,
                max_tokens=8 * len(chunks) + 50,
                messages=[{"role": "user", "content": prompt}]
            )
            if session_id is not None:
                try:
                    await self.db.log_api_call(
                        session_id, "scoring", settings.claude_model,
                        msg.usage.input_tokens, msg.usage.output_tokens
                    )
                except Exception as log_err:
                    logger.warning("Failed to log API usage", error=str(log_err))
            scores = self._parse_batch_scores(msg.content[0].text.strip(), len(chunks))
            logger.debug("Claude batch scored", n=len(chunks),
                         avg=round(sum(scores) / len(scores), 2))
            return scores
        except Exception as e:
            logger.warning("Claude batch scoring failed, will fallback",
                           error=str(e), n=len(chunks))
            return None

    async def _score_with_claude(self, chunk: str, topic: str,
                                  session_id: int | None = None) -> float:
        from src.llm import get_anthropic_client
        prompt = (
            f'Rate 0-10 how relevant this text is to the topic "{topic}". '
            f'Be generous — if the text is tangentially related, score 4+. '
            f'Only score below 3 if completely unrelated. '
            f'Reply with only a number.\n\nText:\n{chunk[:500]}'
        )
        try:
            client = get_anthropic_client()
            msg = await client.messages.create(
                model=settings.claude_model,
                max_tokens=10,
                messages=[{"role": "user", "content": prompt}]
            )
            if session_id is not None:
                try:
                    await self.db.log_api_call(
                        session_id, "scoring", settings.claude_model,
                        msg.usage.input_tokens, msg.usage.output_tokens
                    )
                except Exception as log_err:
                    logger.warning("Failed to log API usage", error=str(log_err))
            response = msg.content[0].text.strip()
            numbers = re.findall(r'\b(\d+(?:\.\d+)?)\b', response)
            if numbers:
                score = float(numbers[0])
                normalized = min(1.0, score / 10.0)
                logger.debug("Claude relevance score", raw=score, normalized=round(normalized, 2))
                return normalized
            return 0.5
        except Exception as e:
            logger.warning("Claude scoring failed, falling back to Ollama", error=str(e))
            return await self._score_with_ollama(chunk, topic)

    async def _score_with_ollama(self, chunk: str, topic: str) -> float:
        prompt = (
            f'Score 0-10: how relevant is this text to the topic "{topic}"?\n'
            f"0 = completely unrelated, 10 = directly and specifically about this topic.\n\n"
            f"Text:\n{chunk[:500]}\n\n"
            f"Reply with ONLY a single integer 0-10. No explanation."
        )
        try:
            response = await self.ollama.generate(prompt, temperature=0.1)
            numbers = re.findall(r'\b(\d+(?:\.\d+)?)\b', response)
            if numbers:
                score = float(numbers[0])
                normalized = min(1.0, score / 10.0)
                logger.debug("Ollama relevance score", raw=score, normalized=round(normalized, 2))
                return normalized
            logger.debug("No number in Ollama relevance response", response=response[:80])
            return 0.6
        except Exception as e:
            logger.warning("Ollama relevance scoring failed", error=str(e))
            return 0.6

    async def rag_query(self, session_id: int, query: str, top_k: int = 20) -> str:
        """
        Retrieve most relevant chunks for a query using embeddings + keyword fallback
        """
        # Get query embedding
        query_embedding = await self.ollama.embed(query)

        # Get top quality chunks
        chunks = await self.db.get_top_chunks(session_id, limit=300)

        if query_embedding and chunks:
            # Rank by embedding similarity
            scored = []
            for chunk in chunks:
                emb = chunk.get("embedding")
                if emb and isinstance(emb, str):
                    try:
                        emb = json.loads(emb)
                    except Exception:
                        emb = None
                sim = cosine_similarity(query_embedding, emb) if emb else 0.5
                scored.append((sim * 0.7 + chunk["quality_score"] * 0.3, chunk))

            scored.sort(key=lambda x: x[0], reverse=True)
            top_chunks = [c for _, c in scored[:top_k]]
        else:
            # Fallback: just use quality score
            top_chunks = chunks[:top_k]

        # Build context
        context_parts = []
        for chunk in top_chunks:
            source_label = f"[{chunk.get('source_type', 'web').upper()}] {chunk.get('title', 'Unknown')}"
            context_parts.append(f"{source_label}:\n{chunk['content']}")

        return "\n\n---\n\n".join(context_parts)