from __future__ import annotations

import asyncio
import math
import os
import random
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Set, Tuple

import pandas as pd

from gabriel.core.prompt_template import PromptTemplate, resolve_template
from gabriel.tasks.deduplicate import Deduplicate, DeduplicateConfig
from gabriel.utils import safest_json
from gabriel.utils.openai_utils import get_all_responses
from gabriel.utils.logging import announce_prompt_rendering


@dataclass
class SeedConfig:
    """Configuration options for :class:`Seed`."""

    instructions: str
    save_dir: str = os.path.expanduser("~/Documents/runs")
    file_name: str = "seed_entities.csv"
    model: str = "gpt-5.2"
    n_parallels: int = 650
    num_entities: int = 1000
    entities_per_generation: int = 50
    entity_batch_frac: float = 0.25
    existing_entities_cap: int = 100
    use_dummy: bool = False
    deduplicate: bool = False
    deduplicate_sample_seed: int = 42
    reasoning_effort: Optional[str] = None


class Seed:
    """Generate structured entity seeds via batched language-model calls."""

    def __init__(
        self,
        cfg: SeedConfig,
        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
        if cfg.num_entities <= 0:
            raise ValueError("num_entities must be positive")
        if cfg.entities_per_generation <= 0:
            raise ValueError("entities_per_generation must be positive")
        if not 0 < cfg.entity_batch_frac <= 1:
            raise ValueError("entity_batch_frac must be between 0 and 1")
        if cfg.existing_entities_cap < 0:
            raise ValueError("existing_entities_cap must be non-negative")
        self.template = resolve_template(
            template=template,
            template_path=template_path,
            reference_filename="seed.jinja2",
        )

    async def run(
        self,
        *,
        existing_entities: Optional[Sequence[str]] = None,
        reset_files: bool = False,
        **response_kwargs: Any,
    ) -> pd.DataFrame:
        """Generate ``num_entities`` unique seed entities."""
        if self.cfg.deduplicate:
            return await self._run_with_deduplication(
                existing_entities=existing_entities,
                reset_files=reset_files,
                **response_kwargs,
            )
        return await self._run_standard(
            existing_entities=existing_entities,
            reset_files=reset_files,
            **response_kwargs,
        )

    async def _run_standard(
        self,
        *,
        existing_entities: Optional[Sequence[str]],
        reset_files: bool,
        **response_kwargs: Any,
    ) -> pd.DataFrame:
        normalized_existing = self._prepare_initial_existing(existing_entities)
        seen: Dict[str, str] = {}
        for ent in normalized_existing:
            norm = self._normalize_entity(ent)
            if norm and norm not in seen:
                seen[norm] = ent

        batch_target = max(
            self.cfg.entities_per_generation,
            math.ceil(self.cfg.num_entities * self.cfg.entity_batch_frac),
        )
        raw_save = os.path.join(self.cfg.save_dir, "seed_raw_responses.csv")
        batch_index = 0
        request_index = 0
        reset_next = reset_files
        while len(seen) < self.cfg.num_entities:
            remaining = self.cfg.num_entities - len(seen)
            batch_goal = min(batch_target, remaining)
            batch_goal = max(batch_goal, self.cfg.entities_per_generation)
            target_count = min(self.cfg.num_entities, len(seen) + batch_goal)
            batch_added = 0
            while len(seen) < target_count:
                remaining_in_batch = target_count - len(seen)
                current_goal = max(remaining_in_batch, self.cfg.entities_per_generation)
                prompts, identifiers = self._build_prompts(
                    current_goal,
                    request_index,
                    list(seen.values()),
                )
                if not prompts:
                    break
                print(
                    f"[Seed] Requesting {len(prompts)} prompts (batch {batch_index}, "
                    f"targeting {current_goal} entities; batch target {batch_goal})."
                )
                df_resp = await self._request_entities(
                    prompts,
                    identifiers,
                    raw_save=raw_save,
                    reset_files=reset_next,
                    **response_kwargs,
                )
                resp_lookup = dict(zip(df_resp.Identifier, df_resp.Response))
                parsed = await asyncio.gather(
                    *[
                        self._parse_entities(resp_lookup.get(identifier, ""))
                        for identifier in identifiers
                    ]
                )
                added = 0
                for entity_list in parsed:
                    for entity in entity_list:
                        norm = self._normalize_entity(entity)
                        if not norm or norm in seen:
                            continue
                        seen[norm] = entity
                        added += 1
                batch_added += added
                reset_next = False
                request_index += 1
                if added == 0 and not any(parsed):
                    break
            print(
                f"[Seed] Added {batch_added} new entities in batch {batch_index}. Total so far: {len(seen)}."
            )
            batch_index += 1
            if batch_added == 0:
                break

        ordered = [seen[norm] for norm in seen]
        trimmed = ordered[: self.cfg.num_entities]
        return self._finalize_entities(trimmed)

    async def _run_with_deduplication(
        self,
        *,
        existing_entities: Optional[Sequence[str]],
        reset_files: bool,
        **response_kwargs: Any,
    ) -> pd.DataFrame:
        normalized_existing = self._prepare_initial_existing(existing_entities)
        all_entities: List[str] = list(normalized_existing)
        seen_norm: Set[str] = set()
        for ent in all_entities:
            norm = self._normalize_entity(ent)
            if norm:
                seen_norm.add(norm)

        dedup_cycle = 0
        deduped = await self._deduplicate_entities(
            all_entities,
            cycle_index=dedup_cycle,
            reset_files=reset_files,
            response_kwargs=response_kwargs,
        )
        dedup_cycle += 1
        raw_target = self.cfg.num_entities * 4
        batch_target = max(
            self.cfg.entities_per_generation,
            math.ceil(raw_target * self.cfg.entity_batch_frac),
        )
        raw_save = os.path.join(self.cfg.save_dir, "seed_raw_responses.csv")
        batch_index = 0
        request_index = 0
        reset_next = reset_files
        while len(deduped) < self.cfg.num_entities:
            cycle_target = max(raw_target, len(all_entities))
            while len(all_entities) < cycle_target:
                remaining_in_cycle = cycle_target - len(all_entities)
                current_goal = min(batch_target, remaining_in_cycle)
                current_goal = max(current_goal, self.cfg.entities_per_generation)
                prompts, identifiers = self._build_prompts(
                    current_goal,
                    request_index,
                    deduped,
                )
                if not prompts:
                    break
                print(
                    f"[Seed] Requesting {len(prompts)} prompts (batch {batch_index}, "
                    f"targeting {current_goal} entities before deduplication)."
                )
                df_resp = await self._request_entities(
                    prompts,
                    identifiers,
                    raw_save=raw_save,
                    reset_files=reset_next,
                    **response_kwargs,
                )
                resp_lookup = dict(zip(df_resp.Identifier, df_resp.Response))
                parsed = await asyncio.gather(
                    *[
                        self._parse_entities(resp_lookup.get(identifier, ""))
                        for identifier in identifiers
                    ]
                )
                added = 0
                for entity_list in parsed:
                    for entity in entity_list:
                        norm = self._normalize_entity(entity)
                        if not norm or norm in seen_norm:
                            continue
                        all_entities.append(entity)
                        seen_norm.add(norm)
                        added += 1
                reset_next = False
                request_index += 1
                if added == 0 and not any(parsed):
                    break

            deduped = await self._deduplicate_entities(
                all_entities,
                cycle_index=dedup_cycle,
                reset_files=False,
                response_kwargs=response_kwargs,
            )
            dedup_cycle += 1
            print(
                f"[Seed] Unique after deduplication: {len(deduped)}."
            )
            if len(deduped) >= self.cfg.num_entities:
                break

            all_entities = list(deduped)
            seen_norm = {
                self._normalize_entity(entity)
                for entity in all_entities
                if self._normalize_entity(entity)
            }
            batch_index += 1
            if not all_entities:
                break

        trimmed = self._sample_to_target(deduped)
        return self._finalize_entities(trimmed)

    async def _request_entities(
        self,
        prompts: List[str],
        identifiers: List[str],
        *,
        raw_save: str,
        reset_files: bool,
        **response_kwargs: Any,
    ) -> pd.DataFrame:
        kwargs = dict(response_kwargs)
        kwargs.setdefault("model", self.cfg.model)
        kwargs.setdefault("n_parallels", self.cfg.n_parallels)
        kwargs.setdefault("use_dummy", self.cfg.use_dummy)
        kwargs.setdefault("reasoning_effort", self.cfg.reasoning_effort)
        kwargs.setdefault("json_mode", True)
        kwargs.setdefault("save_path", raw_save)
        kwargs.setdefault("reset_files", reset_files)
        df_resp = await get_all_responses(
            prompts=prompts,
            identifiers=identifiers,
            **kwargs,
        )
        if not isinstance(df_resp, pd.DataFrame):
            raise RuntimeError("get_all_responses returned no DataFrame")
        return df_resp

    def _finalize_entities(self, entities: List[str]) -> pd.DataFrame:
        df = pd.DataFrame(
            {
                "entity": entities,
                "entity_id": [f"entity-{idx:05d}" for idx in range(len(entities))],
            }
        )
        df["source_batch"] = df.index // max(self.cfg.entities_per_generation, 1)
        df["source_identifier"] = ["seed" for _ in range(len(entities))]
        final_path = os.path.join(self.cfg.save_dir, self.cfg.file_name)
        df.to_csv(final_path, index=False)
        print(
            f"[Seed] Generated {len(df)} entities. Saved aggregated seeds to {final_path}."
        )
        return df

    async def _parse_entities(self, raw: Any) -> List[str]:
        obj = await safest_json(raw)
        results: List[str] = []
        if isinstance(obj, dict):
            for key in sorted(obj.keys()):
                value = obj.get(key)
                if value is None:
                    continue
                text = str(value).strip()
                if text:
                    results.append(text)
        elif isinstance(obj, list):
            for item in obj:
                if item is None:
                    continue
                text = str(item).strip()
                if text:
                    results.append(text)
        elif isinstance(obj, str):
            text = obj.strip()
            if text:
                results.append(text)
        return results

    def _prepare_initial_existing(
        self, entries: Optional[Sequence[str]]
    ) -> List[str]:
        if not entries:
            return []
        unique: List[str] = []
        seen: Set[str] = set()
        for entry in entries:
            text = str(entry).strip()
            if not text:
                continue
            norm = self._normalize_entity(text)
            if norm and norm not in seen:
                seen.add(norm)
                unique.append(text)
        return unique

    def _build_prompts(
        self,
        goal: int,
        batch_index: int,
        seen_entities: Sequence[str],
    ) -> Tuple[List[str], List[str]]:
        prompts: List[str] = []
        identifiers: List[str] = []
        per_call = self.cfg.entities_per_generation
        prompt_count = math.ceil(goal / per_call)
        existing_sample = self._sample_existing(seen_entities)
        existing_blob = "\n".join(existing_sample) if existing_sample else None
        announce_prompt_rendering("Seed", prompt_count)
        for call_index in range(prompt_count):
            identifiers.append(f"seed|{batch_index}|{call_index}")
            prompts.append(
                self.template.render(
                    instructions=self.cfg.instructions,
                    entities_per_generation=per_call,
                    existing_entities=existing_blob,
                )
            )
        return prompts, identifiers

    def _sample_existing(self, seen_entities: Sequence[str]) -> List[str]:
        if not seen_entities:
            return []
        cap = max(0, self.cfg.existing_entities_cap)
        if cap == 0:
            return []
        if len(seen_entities) <= cap:
            return list(seen_entities)
        return random.sample(list(seen_entities), cap)

    async def _deduplicate_entities(
        self,
        entities: Sequence[str],
        *,
        cycle_index: int,
        reset_files: bool,
        response_kwargs: Dict[str, Any],
    ) -> List[str]:
        cleaned: List[str] = []
        for entity in entities:
            text = str(entity).strip()
            if text:
                cleaned.append(text)
        if not cleaned:
            return []

        df = pd.DataFrame({"entity": cleaned})
        dedup_cfg = DeduplicateConfig(
            save_dir=os.path.join(self.cfg.save_dir, "seed_deduplicate"),
            file_name=f"seed_deduplicate_cycle{cycle_index}.csv",
            model=self.cfg.model,
            n_parallels=self.cfg.n_parallels,
            n_runs=4,
            use_dummy=self.cfg.use_dummy,
            group_size=100,
        )
        dedup = Deduplicate(dedup_cfg)
        dedup_df = await dedup.run(
            df,
            column_name="entity",
            reset_files=reset_files,
            **self._filter_dedup_response_kwargs(response_kwargs),
        )
        mapped_col = (
            "mapped_entity_final"
            if "mapped_entity_final" in dedup_df.columns
            else "mapped_entity"
        )
        unique: List[str] = []
        seen: Set[str] = set()
        for val in dedup_df[mapped_col]:
            text = str(val).strip()
            if not text or text in seen:
                continue
            seen.add(text)
            unique.append(text)
        return unique

    @staticmethod
    def _filter_dedup_response_kwargs(response_kwargs: Dict[str, Any]) -> Dict[str, Any]:
        blocked = {
            "model",
            "n_parallels",
            "save_path",
            "json_mode",
            "reset_files",
            "prompts",
            "identifiers",
        }
        return {key: value for key, value in response_kwargs.items() if key not in blocked}

    def _sample_to_target(self, entities: List[str]) -> List[str]:
        if len(entities) <= self.cfg.num_entities:
            return entities[: self.cfg.num_entities]
        rng = random.Random(self.cfg.deduplicate_sample_seed)
        selected = set(rng.sample(entities, self.cfg.num_entities))
        return [entity for entity in entities if entity in selected][: self.cfg.num_entities]

    @staticmethod
    def _normalize_entity(text: str) -> str:
        collapsed = re.sub(r"\s+", " ", text.strip()).lower()
        return collapsed
