import os
import re
import time
import pandas as pd
from typing import Optional, Dict, List

# Lazy imports so users don't need both SDKs installed to import the module
def _lazy_import_groq():
    from groq import Groq
    return Groq

def _lazy_import_openai():
    from openai import OpenAI
    return OpenAI

class BioEvaluator:
    """
    Cross-client evaluator that can call Groq (default) or OpenAI chat completions.
    Expects a DataFrame with a 'Question' column.
    Produces one CSV per question and a summary.csv.
    """
    def __init__(self, model_name: str, client: str = "groq", api_key: Optional[str] = None):
        self.model_name = model_name
        self.client_name = client.lower()
        self.api_key = api_key
        self._client = None

        if self.client_name not in {"groq", "openai"}:
            raise ValueError("client must be 'groq' or 'openai'.")

        # Resolve API key from env if not provided explicitly
        if not self.api_key:
            env_key = "GROQ_API_KEY" if self.client_name == "groq" else "OPENAI_API_KEY"
            self.api_key = os.getenv(env_key)

        if not self.api_key:
            raise ValueError(
                "API key not provided. Pass api_key=... or set GROQ_API_KEY / OPENAI_API_KEY."
            )

        # Initialize SDK client
        if self.client_name == "groq":
            Groq = _lazy_import_groq()
            self._client = Groq(api_key=self.api_key)
        else:
            OpenAI = _lazy_import_openai()
            self._client = OpenAI(api_key=self.api_key)

    # ---------- helpers ----------
    def _extract_table(self, text: str) -> pd.DataFrame:
        """
        Extract the first markdown-style table from the text into a DataFrame.
        If none exists, fall back to a 1-column DataFrame with the raw answer.
        """
        tables = re.findall(r'((?:\|.*\n)+)', text or "", flags=re.DOTALL)
        if not tables:
            return pd.DataFrame([{"Answer": (text or "").strip()}])

        table_text = tables[0].strip()
        lines = [ln for ln in (l.strip() for l in table_text.split("\n")) if ln.startswith("|")]
        table = [ln.strip("|").split("|") for ln in lines]

        # Remove header separators like |---|:---:|
        table = [row for row in table if not all(set(c.strip()) <= {"-", ":"} for c in row)]
        if len(table) < 2:
            return pd.DataFrame([{"Answer": (text or "").strip()}])

        header = [h.strip() for h in table[0]]
        rows = [[c.strip() for c in r] for r in table[1:]]
        rows = [r if len(r) == len(header) else r + [""] * (len(header) - len(r)) for r in rows]
        return pd.DataFrame(rows, columns=header)

    @staticmethod
    def _safe_filename(text: str, maxlen: int = 80) -> str:
        cleaned = "".join(ch if ch.isalnum() or ch in (" ", "_", "-") else "_" for ch in text).strip()
        cleaned = re.sub(r"\\s+", "_", cleaned)
        return (cleaned[:maxlen] or "question").rstrip("_")

    def _chat_completion(self, messages: List[Dict[str, str]]) -> str:
        """
        Provider-agnostic chat call; returns text content.
        """
        if self.client_name == "groq":
            completion = self._client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                temperature=0.0,
                max_completion_tokens=4000,
                top_p=1,
                stream=False,
            )
            return completion.choices[0].message.content or ""

        # OpenAI
        completion = self._client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            temperature=0.0,
            max_tokens=4000,
            top_p=1,
            stream=False,
        )
        return completion.choices[0].message.content or ""

    # ---------- main ----------
    def evaluate(self, df: pd.DataFrame, output_folder: str = "answers_per_question") -> str:
        """
        Create one CSV per question and a summary.csv with (csv_file, num_rows, seconds).
        Returns the path to summary.csv.
        """
        if "Question" not in df.columns:
            raise ValueError("Input DataFrame must contain a 'Question' column.")

        os.makedirs(output_folder, exist_ok=True)
        summary_rows = []

        prompt_instruction = (
            "Respond ONLY in a tabular format, sourced from a reputable biological database or repository "
            "— not blogs or Wikipedia. Include a citation column with the source for each row. "
            'Each cell must contain a single atomic value (no lists, commas, or separators).'
        )

        used_filenames = set()

        for _, row in df.iterrows():
            question = str(row["Question"]).strip()
            start = time.perf_counter()
            num_rows = 0
            csv_file = "ERROR"

            try:
                messages = [
                    {"role": "system", "content": "You are a biological domain expert."},
                    {"role": "user", "content": f"{prompt_instruction}\\n\\nQuestion: {question}"},
                ]

                answer = self._chat_completion(messages)
                table_df = self._extract_table(answer)

                # unique filename
                base = self._safe_filename(question)
                filename = base
                i = 2
                while filename in used_filenames:
                    filename = f"{base}_{i}"
                    i += 1
                used_filenames.add(filename)

                csv_file = f"{filename}.csv"
                out_path = os.path.join(output_folder, csv_file)
                table_df.to_csv(out_path, index=False)
                num_rows = int(table_df.shape[0])
                print(f"Saved: {out_path}")

            except Exception as e:
                print(f"Error for question '{question}': {e}")

            finally:
                duration = time.perf_counter() - start
                summary_rows.append({
                    "csv_file": csv_file,
                    "num_rows": num_rows,
                    "seconds": round(duration, 3),
                })

        summary_df = pd.DataFrame(summary_rows, columns=["csv_file", "num_rows", "seconds"])
        summary_path = os.path.join(output_folder, "summary.csv")
        summary_df.to_csv(summary_path, index=False)
        print(f"\\nSummary written to: {summary_path}")
        return summary_path


# -------- functional wrapper ----------
def generate_response(
    df: pd.DataFrame,
    client: str = "groq",
    model: str = "gemma2-9b-it",
    output_folder: str = "answers_per_question",
    api_key: Optional[str] = None,
):
    """
    Runs evaluation and returns (summary_df, summary_csv_path)
    """
    evaluator = BioEvaluator(model_name=model, client=client, api_key=api_key)
    summary_csv_path = evaluator.evaluate(df, output_folder=output_folder)
    import pandas as pd
    summary_df = pd.read_csv(summary_csv_path)
    return summary_df, summary_csv_path
