from typing import Tuple, Dict, List
import gzip
import json
import numpy as np
from functools import reduce

from blueness import module
from bluer_options.logger.config import log_list
from bluer_objects import file, storage
from bluer_objects import objects
from bluer_objects.html_report import HTMLReport
from bluer_objects.storage.policies import DownloadPolicy

from bluer_agent import NAME
from bluer_agent.host import signature
from bluer_agent.rag.corpus.embed import embed_fn
from bluer_agent.logger import logger

NAME = module.name(__file__, NAME)


def _cosine(a: np.ndarray, b: np.ndarray) -> float:
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))


class Context:
    def __init__(
        self,
        object_name: str,
        download: bool = False,
    ):
        self.list_of_roots: List[str] = []

        self.object_name = object_name

        if download:
            storage.download(
                self.object_name,
                policy=DownloadPolicy.DOESNT_EXIST,
            )

        # --- load corpus embeddings ---
        self.corpus_vec = np.load(
            objects.path_of(
                object_name=object_name,
                filename="corpus.embeddings.npy",
            )
        )

        self.corpus_meta_file = objects.path_of(
            object_name=self.object_name,
            filename="corpus.meta.jsonl.gz",
        )

        self.corpus_text_file = objects.path_of(
            object_name=self.object_name,
            filename="corpus.jsonl.gz",
        )

        logger.info(
            "{} created for {}.".format(
                self.__class__.__name__,
                self.object_name,
            )
        )

    def generate(
        self,
        query: str,
        top_k: int = 5,
        html_report: HTMLReport = HTMLReport(),
    ) -> Tuple[bool, Dict]:
        logger.info(
            '{}[{}].generate("{}")'.format(
                self.__class__.__name__,
                self.object_name,
                query,
            )
        )

        # --- embed query ---
        q_vec = np.asarray(embed_fn([query])[0], dtype=np.float32)

        candidates: List[Tuple[float, dict]] = []

        self.list_of_roots = []
        with gzip.open(
            self.corpus_meta_file,
            "rt",
            encoding="utf-8",
        ) as meta_f, gzip.open(
            self.corpus_text_file,
            "rt",
            encoding="utf-8",
        ) as text_f:
            for i, (meta_line, text_line) in enumerate(zip(meta_f, text_f)):
                meta = json.loads(meta_line)

                if meta["root"] not in self.list_of_roots:
                    self.list_of_roots.append(meta["root"])

                record = json.loads(text_line)
                score = _cosine(q_vec, self.corpus_vec[i])

                candidates.append(
                    (
                        score,
                        {
                            "root": meta["root"],
                            "url": meta["url"],
                            "chunk_id": meta["chunk_id"],
                            "text": record["text"],
                        },
                    )
                )

        candidates.sort(key=lambda x: x[0], reverse=True)
        top = candidates[:top_k]

        log_list(logger, "reviewed", self.list_of_roots, "root(s)")

        chunks = [
            {
                **item,
                "score": round(score, 4),
            }
            for score, item in top
        ]

        for chunk in chunks:
            logger.info(
                "{} #{} @ {:.2f}: {}".format(
                    chunk["root"],
                    chunk["chunk_id"],
                    chunk["score"],
                    chunk["text"][:300],
                )
            )

        context = {
            "chunks": chunks,
        }

        # TODO: change to HTMLReport - 2026-02-05 - nice to have.
        file.save_text(
            objects.path_of(
                object_name=self.object_name,
                filename="result.txt",
            ),
            [
                f"query: {query}",
                "",
                "retrievals:",
            ]
            + [
                "{} #{} @ {:.2f}: {}".format(
                    chunk["root"],
                    chunk["chunk_id"],
                    chunk["score"],
                    chunk["text"][:300],
                )
                for chunk in chunks
            ],
            log=True,
        )

        html_report.replace(
            {
                "context_count:::": str(len(context["chunks"])),
                "corpus_name:::": self.object_name,
                "query:::": query,
                "query_dir:::": "ltr",
                "query_lang:::": "fa",
                "host_signature:::": " | ".join(signature()),
                "text_dir:::": "ltr",
                "text_lang:::": "fa",
                "title:::": "rag query output",
            }
        ).replace(
            {
                "context:::": reduce(
                    lambda x, y: x + y,
                    [
                        [
                            "<tr>",
                            '    <td dir="ltr">',
                            '        <a href="{}" target="_blank" rel="noreferrer" title="{}">{}#{}</a>'.format(
                                chunk["url"],
                                chunk["text"],
                                chunk["root"],
                                chunk["chunk_id"],
                            ),
                            "    </td>",
                            '    <td class="mono">{:.2f}</td>'.format(chunk["score"]),
                            "</tr>",
                            "",
                        ]
                        for chunk in context["chunks"]
                    ],
                    [],
                ),
            },
            contains=True,
        )

        return True, context

    def understand_reply(self, reply: str) -> Tuple[bool, str]:
        if len(self.list_of_roots) == 1:
            return True, reply

        try:
            reply_dict = json.loads(reply)
        except Exception as e:
            logger.error(e)
            return False, ""

        for field in [
            "winner_root",
            "argument_fa",
        ]:
            if field not in reply_dict:
                logger.error(f"{field} not found.")
                return False, ""

        logger.info("🥇 winner is {}!".format(reply_dict["winner_root"]))

        return True, reply_dict["argument_fa"]
