from __future__ import annotations

import asyncio
import hashlib
import os
import random
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional

import pandas as pd

from ..core.prompt_template import PromptTemplate, resolve_template
from ..utils.openai_utils import get_all_responses
from ..utils import (
    safest_json,
    load_image_inputs,
    load_audio_inputs,
    load_pdf_inputs,
    warn_if_modality_mismatch,
)
from ..utils.logging import announce_prompt_rendering


@dataclass
class CompareConfig:
    save_dir: str = "comparison"
    file_name: str = "comparison_responses.csv"
    model: str = "gpt-5-mini"
    n_parallels: int = 650
    n_runs: int = 1
    use_dummy: bool = False
    differentiate: bool = True
    additional_instructions: Optional[str] = None
    modality: str = "text"
    reasoning_effort: Optional[str] = None
    circle_first: Optional[bool] = None

    def __post_init__(self) -> None:
        if self.additional_instructions is not None:
            cleaned = str(self.additional_instructions).strip()
            self.additional_instructions = cleaned or None


class Compare:
    """Compare two columns row-wise using an LLM."""

    def __init__(
        self,
        cfg: CompareConfig,
        template: Optional[PromptTemplate] = None,
        template_path: Optional[str] = None,
    ) -> None:
        expanded = Path(os.path.expandvars(os.path.expanduser(cfg.save_dir)))
        expanded.mkdir(parents=True, exist_ok=True)
        cfg.save_dir = str(expanded)
        self.cfg = cfg
        self.template = resolve_template(
            template=template,
            template_path=template_path,
            reference_filename="comparison_prompt.jinja2",
        )

    async def _parse(self, raw: Any) -> Dict[str, str]:
        obj = await safest_json(raw)
        if isinstance(obj, dict):
            return {str(k): str(v) if v is not None else "" for k, v in obj.items()}
        return {}

    async def run(
        self,
        df: pd.DataFrame,
        circle_column_name: str,
        square_column_name: str,
        *,
        reset_files: bool = False,
        **kwargs: Any,
    ) -> pd.DataFrame:
        df_proc = df.reset_index(drop=True).copy()
        mask = df_proc[circle_column_name].notna() & df_proc[square_column_name].notna()
        skipped = int((~mask).sum())
        if skipped:
            print(
                f"Skipping {skipped} rows with NaN in {circle_column_name} or {square_column_name}"
            )
        df_proc = df_proc[mask].reset_index(drop=True)

        circles = df_proc[circle_column_name].tolist()
        squares = df_proc[square_column_name].tolist()
        pairs = list(zip(circles, squares))

        warn_if_modality_mismatch(circles, self.cfg.modality, column_name=circle_column_name)
        warn_if_modality_mismatch(squares, self.cfg.modality, column_name=square_column_name)

        prompts: List[str] = []
        ids: List[str] = []
        id_to_circle_first: Dict[str, bool] = {}
        prompt_circle_text: Dict[str, str] = {}
        prompt_square_text: Dict[str, str] = {}
        for circle, square in pairs:
            ident = hashlib.sha1(f"{circle}|{square}".encode()).hexdigest()[:8]
            ids.append(ident)
            circle_first_flag = (
                self.cfg.circle_first
                if self.cfg.circle_first is not None
                else random.random() < 0.5
            )
            id_to_circle_first[ident] = circle_first_flag
            prompt_circle_text[ident] = (
                circle if self.cfg.modality in {"text", "entity", "web"} else ""
            )
            prompt_square_text[ident] = (
                square if self.cfg.modality in {"text", "entity", "web"} else ""
            )

        announce_prompt_rendering("Compare", len(ids))

        for ident in ids:
            prompts.append(
                self.template.render(
                    entry_circle=prompt_circle_text[ident],
                    entry_square=prompt_square_text[ident],
                    differentiate=self.cfg.differentiate,
                    additional_instructions=self.cfg.additional_instructions or "",
                    modality=self.cfg.modality,
                    circle_first=id_to_circle_first[ident],
                )
            )

        prompt_images: Optional[Dict[str, List[str]]] = None
        prompt_audio: Optional[Dict[str, List[Dict[str, str]]]] = None
        prompt_pdfs: Optional[Dict[str, List[Dict[str, str]]]] = None
        if self.cfg.modality == "image":
            tmp: Dict[str, List[str]] = {}
            for ident, (circle, square) in zip(ids, pairs):
                imgs: List[str] = []
                circle_imgs = load_image_inputs(circle)
                square_imgs = load_image_inputs(square)
                if id_to_circle_first.get(ident, False):
                    if circle_imgs:
                        imgs.extend(circle_imgs)
                    if square_imgs:
                        imgs.extend(square_imgs)
                else:
                    if square_imgs:
                        imgs.extend(square_imgs)
                    if circle_imgs:
                        imgs.extend(circle_imgs)
                if imgs:
                    tmp[ident] = imgs
            prompt_images = tmp or None
        elif self.cfg.modality == "audio":
            tmp_a: Dict[str, List[Dict[str, str]]] = {}
            for ident, (circle, square) in zip(ids, pairs):
                auds: List[Dict[str, str]] = []
                circle_auds = load_audio_inputs(circle)
                square_auds = load_audio_inputs(square)
                if id_to_circle_first.get(ident, False):
                    if circle_auds:
                        auds.extend(circle_auds)
                    if square_auds:
                        auds.extend(square_auds)
                else:
                    if square_auds:
                        auds.extend(square_auds)
                    if circle_auds:
                        auds.extend(circle_auds)
                if auds:
                    tmp_a[ident] = auds
            prompt_audio = tmp_a or None
        elif self.cfg.modality == "pdf":
            tmp_p: Dict[str, List[Dict[str, str]]] = {}
            for ident, (circle, square) in zip(ids, pairs):
                pdfs: List[Dict[str, str]] = []
                circle_pdfs = load_pdf_inputs(circle)
                square_pdfs = load_pdf_inputs(square)
                if id_to_circle_first.get(ident, False):
                    if circle_pdfs:
                        pdfs.extend(circle_pdfs)
                    if square_pdfs:
                        pdfs.extend(square_pdfs)
                else:
                    if square_pdfs:
                        pdfs.extend(square_pdfs)
                    if circle_pdfs:
                        pdfs.extend(circle_pdfs)
                if pdfs:
                    tmp_p[ident] = pdfs
            prompt_pdfs = tmp_p or None

        csv_path = os.path.join(self.cfg.save_dir, self.cfg.file_name)

        kwargs.setdefault("web_search", self.cfg.modality == "web")

        df_resp_all = await get_all_responses(
            prompts=prompts,
            identifiers=ids,
            prompt_images=prompt_images,
            prompt_audio=prompt_audio,
            prompt_pdfs=prompt_pdfs,
            n_parallels=self.cfg.n_parallels,
            model=self.cfg.model,
            save_path=csv_path,
            use_dummy=self.cfg.use_dummy,
            json_mode=self.cfg.modality != "audio",
            reset_files=reset_files,
            reasoning_effort=self.cfg.reasoning_effort,
            **kwargs,
        )
        if not isinstance(df_resp_all, pd.DataFrame):
            raise RuntimeError("get_all_responses returned no DataFrame")

        resp_map = dict(zip(df_resp_all.Identifier, df_resp_all.Response))
        parsed = await asyncio.gather(*[self._parse(resp_map.get(i, "")) for i in ids])

        records: List[Dict[str, str]] = []
        for (circle, square), res in zip(pairs, parsed):
            for attr, expl in res.items():
                records.append(
                    {
                        circle_column_name: circle,
                        square_column_name: square,
                        "attribute": attr,
                        "explanation": expl,
                    }
                )

        out_df = pd.DataFrame(records)
        if not out_df.empty:
            out_df.set_index([circle_column_name, square_column_name], inplace=True)
        return out_df
