"""
Compare — side-by-side comparison of two prompt versions.

You bring the responses. PromptVault shows you the difference.
"""

from typing import Any, Dict, List, Optional

from .prompt import Prompt
from .storage import PromptStorage
from pathlib import Path


class Compare:
    """
    Compare two prompt versions side by side.

    You call your LLM for both prompts. You hand the responses
    to PromptVault. It logs, stores, and summarises the comparison.

    Example:
        >>> v1 = Prompt("summarize-v1", "Summarize this: {text}")
        >>> v2 = Prompt("summarize-v2", "Summarize in 3 bullets: {text}")
        >>>
        >>> # You call your own LLM
        >>> r1 = my_llm(v1.render(text="Some article..."))
        >>> r2 = my_llm(v2.render(text="Some article..."))
        >>>
        >>> # PromptVault compares them
        >>> cmp = Compare(v1, v2)
        >>> cmp.log(r1, r2)
        >>> cmp.show()
    """

    def __init__(self, prompt_a: Prompt, prompt_b: Prompt):
        self.prompt_a = prompt_a
        self.prompt_b = prompt_b
        self._runs: List[Dict[str, Any]] = []

        # Share the same vault dir as the prompts
        self._storage = PromptStorage(prompt_a._vault_dir)

    def log(
        self,
        response_a: str,
        response_b: str,
        rendered_a: Optional[str] = None,
        rendered_b: Optional[str] = None,
        model: Optional[str] = None,
        latency_ms_a: Optional[float] = None,
        latency_ms_b: Optional[float] = None,
        tokens_a: Optional[int] = None,
        tokens_b: Optional[int] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, str]:
        """
        Log a comparison run — one response from each prompt variant.

        Args:
            response_a:   The LLM response for prompt_a.
            response_b:   The LLM response for prompt_b.
            rendered_a:   The final rendered text sent to the LLM for prompt_a.
            rendered_b:   The final rendered text sent to the LLM for prompt_b.
            model:        The LLM model you used (e.g. 'gpt-4', 'llama3').
            latency_ms_a: Time in ms the LLM took for prompt_a.
            latency_ms_b: Time in ms the LLM took for prompt_b.
            tokens_a:     Token count for prompt_a response.
            tokens_b:     Token count for prompt_b response.
            extra:        Any additional metadata.

        Returns:
            Dict with run_ids for both sides: {'run_id_a': ..., 'run_id_b': ...}
        """
        run_id_a = self._storage.log_run(
            prompt_name=self.prompt_a.name,
            prompt_version=self.prompt_a.version,
            rendered_prompt=rendered_a or "",
            response=response_a,
            model=model,
            latency_ms=latency_ms_a,
            tokens=tokens_a,
            extra={**(extra or {}), "comparison": True},
        )
        run_id_b = self._storage.log_run(
            prompt_name=self.prompt_b.name,
            prompt_version=self.prompt_b.version,
            rendered_prompt=rendered_b or "",
            response=response_b,
            model=model,
            latency_ms=latency_ms_b,
            tokens=tokens_b,
            extra={**(extra or {}), "comparison": True},
        )

        pair = {
            "run_id_a": run_id_a,
            "run_id_b": run_id_b,
            "response_a": response_a,
            "response_b": response_b,
            "model": model,
            "latency_ms_a": latency_ms_a,
            "latency_ms_b": latency_ms_b,
            "tokens_a": tokens_a,
            "tokens_b": tokens_b,
        }
        self._runs.append(pair)
        return {"run_id_a": run_id_a, "run_id_b": run_id_b}

    def diff(self, index: int = -1) -> Dict[str, Any]:
        """
        Return a structured diff of the most recent (or nth) comparison.

        Args:
            index: Which run to diff. Default -1 = most recent.

        Returns:
            Dict with side-by-side details and simple metrics.
        """
        if not self._runs:
            raise RuntimeError("No comparisons logged yet. Call .log() first.")

        run = self._runs[index]
        ra, rb = run["response_a"], run["response_b"]

        words_a = len(ra.split())
        words_b = len(rb.split())
        chars_a = len(ra)
        chars_b = len(rb)

        return {
            "prompt_a": f"{self.prompt_a.name} ({self.prompt_a.version})",
            "prompt_b": f"{self.prompt_b.name} ({self.prompt_b.version})",
            "model":    run.get("model"),
            "response_a":   ra,
            "response_b":   rb,
            "words_a":      words_a,
            "words_b":      words_b,
            "chars_a":      chars_a,
            "chars_b":      chars_b,
            "word_diff":    words_b - words_a,
            "latency_ms_a": run.get("latency_ms_a"),
            "latency_ms_b": run.get("latency_ms_b"),
            "tokens_a":     run.get("tokens_a"),
            "tokens_b":     run.get("tokens_b"),
        }

    def summary(self) -> Dict[str, Any]:
        """
        Aggregate summary across all logged comparison runs.

        Returns:
            Stats for both variants including avg latency, tokens, response length.
        """
        if not self._runs:
            raise RuntimeError("No comparisons logged yet. Call .log() first.")

        def avg(lst):
            clean = [x for x in lst if x is not None]
            return round(sum(clean) / len(clean), 2) if clean else None

        return {
            "prompt_a": f"{self.prompt_a.name} ({self.prompt_a.version})",
            "prompt_b": f"{self.prompt_b.name} ({self.prompt_b.version})",
            "total_comparisons": len(self._runs),
            "avg_words_a":      avg([len(r["response_a"].split()) for r in self._runs]),
            "avg_words_b":      avg([len(r["response_b"].split()) for r in self._runs]),
            "avg_latency_ms_a": avg([r.get("latency_ms_a") for r in self._runs]),
            "avg_latency_ms_b": avg([r.get("latency_ms_b") for r in self._runs]),
            "avg_tokens_a":     avg([r.get("tokens_a") for r in self._runs]),
            "avg_tokens_b":     avg([r.get("tokens_b") for r in self._runs]),
        }

    def show(self, index: int = -1):
        """
        Print a clean side-by-side comparison to the terminal.

        Args:
            index: Which run to show. Default -1 = most recent.
        """
        d = self.diff(index)
        width = 60
        divider = "─" * (width * 2 + 7)

        print(f"\n{'─' * width}")
        print(f"  PROMPTVAULT COMPARISON")
        print(f"{'─' * width}")
        print(f"  {'Prompt A':<30} {d['prompt_a']}")
        print(f"  {'Prompt B':<30} {d['prompt_b']}")
        if d["model"]:
            print(f"  {'Model':<30} {d['model']}")
        print(f"{'─' * width}")

        print(f"\n  ── Response A ──")
        print(f"  {d['response_a']}\n")
        print(f"  ── Response B ──")
        print(f"  {d['response_b']}\n")

        print(f"{'─' * width}")
        print(f"  {'Metric':<28} {'Prompt A':>12} {'Prompt B':>12}")
        print(f"  {'─'*28} {'─'*12} {'─'*12}")
        print(f"  {'Word count':<28} {str(d['words_a']):>12} {str(d['words_b']):>12}")
        print(f"  {'Char count':<28} {str(d['chars_a']):>12} {str(d['chars_b']):>12}")

        if d["latency_ms_a"] is not None:
            print(f"  {'Latency (ms)':<28} {str(d['latency_ms_a']):>12} {str(d['latency_ms_b']):>12}")
        if d["tokens_a"] is not None:
            print(f"  {'Tokens':<28} {str(d['tokens_a']):>12} {str(d['tokens_b']):>12}")

        print(f"{'─' * width}\n")
