"""EvalPipeline: orchestrates a full evaluation run."""

from __future__ import annotations

import json
import time
import uuid

from rich.console import Console
from rich.progress import track

from promptforge.core.prompt_spec import PromptSpec
from promptforge.core.dataset import Dataset, TestCase
from promptforge.core.run_config import RunConfig
from promptforge.core.templating import render_template
from promptforge.llm.client_base import LLMResponse
from promptforge.eval.heuristics import _resolve_heuristic
from promptforge.store.repositories import RunRepository, CaseResultRepository, ScoreRepository


def _build_llm_client(config: RunConfig):
    if config.provider == "openai":
        from promptforge.llm.openai_client import OpenAIClient
        return OpenAIClient(model=config.model, params=config.params)
    elif config.provider == "anthropic":
        from promptforge.llm.anthropic_client import AnthropicClient
        return AnthropicClient(model=config.model, params=config.params)
    else:
        raise ValueError(f"Unknown provider: {config.provider}")


class EvalPipeline:
    def __init__(
        self,
        prompt_spec: PromptSpec,
        dataset: Dataset,
        run_config: RunConfig,
        console: Console | None = None,
    ) -> None:
        self.ps = prompt_spec
        self.ds = dataset
        self.rc = run_config
        self.console = console or Console()
        self.run_repo = RunRepository()
        self.case_repo = CaseResultRepository()
        self.score_repo = ScoreRepository()

    def run(self) -> str:
        run_id = str(uuid.uuid4())

        self.run_repo.create_run(
            run_id=run_id,
            prompt_id=self.ps.id,
            prompt_version=self.ps.version,
            prompt_hash=self.ps.content_hash,
            dataset_id=self.ds.dataset_id,
            dataset_hash=self.ds.content_hash,
            model=self.rc.model,
            provider=self.rc.provider,
            params=self.rc.params,
        )

        client = _build_llm_client(self.rc)
        total_tokens = 0

        for case in track(self.ds.cases, description="Evaluating...", console=self.console):
            rendered = render_template(self.ps.template, case.input)
            t0 = time.time()
            response: LLMResponse = client.complete(
                rendered,
                self.rc.params,
                system_prompt=self.ps.system_prompt or None,
)
            latency_ms = (time.time() - t0) * 1000
            total_tokens += response.total_tokens

            parsed = self._try_parse_json(response.content)

            self.case_repo.save(
                run_id=run_id,
                case_id=case.id,
                input_json=json.dumps(case.input),
                output_raw=response.content,
                output_parsed_json=json.dumps(parsed) if parsed else None,
                latency_ms=latency_ms,
                tokens_in=response.prompt_tokens,
                tokens_out=response.completion_tokens,
            )

            scores = self._evaluate(case, response.content, parsed, rendered)
            for evaluator_name, score, rationale in scores:
                self.score_repo.save(
                    run_id=run_id,
                    case_id=case.id,
                    evaluator=evaluator_name,
                    dimension=evaluator_name,
                    score=score,
                    rationale=rationale,
                )

        self.run_repo.finish_run(run_id, total_tokens=total_tokens)
        return run_id

    def _try_parse_json(self, text: str) -> dict | None:
        clean = text.strip()
        # Remove markdown code blocks que alguns modelos adicionam
        if clean.startswith("```"):
            lines = clean.split("\n")
            clean = "\n".join(lines[1:])
        if clean.endswith("```"):
            lines = clean.split("\n")
            clean = "\n".join(lines[:-1])
        try:
            return json.loads(clean.strip())
        except (json.JSONDecodeError, ValueError):
            return None

    def _evaluate(
        self,
        case: TestCase,
        output_raw: str,
        output_parsed: dict | None,
        rendered_prompt: str = "",
    ) -> list[tuple[str, float, str]]:
        results = []
        for ev_conf in self.rc.evaluators:

            if ev_conf.type == "heuristic":
                fn = _resolve_heuristic(ev_conf.name)
                if fn is None:
                    self.console.print(
                        f"[yellow]⚠ Unknown heuristic evaluator: '{ev_conf.name}'. Skipping.[/yellow]"
                    )
                    continue
                score, rationale = fn(
                    output_raw=output_raw,
                    output_parsed=output_parsed,
                    expected=case.expected,
                    config=ev_conf.config,
                    prompt_spec=self.ps,
                )
                results.append((ev_conf.name, score, rationale))

            elif ev_conf.type == "judge":
                # LLM-as-judge agora ligado ao pipeline
                try:
                    from promptforge.eval.llm_judge import judge_output
                    from promptforge.eval.rubrics import Rubric

                    rubric_path = ev_conf.config.get("rubric")
                    if not rubric_path:
                        self.console.print(
                            f"[yellow]⚠ Judge '{ev_conf.name}' não tem 'rubric' definida no config. Skipping.[/yellow]"
                        )
                        continue

                    rubric = Rubric.from_yaml(rubric_path)
                    judge_scores = judge_output(
                        output=output_raw,
                        input_context=rendered_prompt,
                        rubric=rubric,
                    )

                    # Cada dimensão da rubrica gera um score separado
                    for dimension, result in judge_scores.items():
                        if isinstance(result, dict):
                            raw_score = result.get("score", 0)
                            rationale = result.get("rationale", "")
                        else:
                            # Fallback: o modelo devolveu só o número
                            raw_score = result
                            rationale = ""

                        # Normaliza o score para 0.0-1.0
                        dim_def = next((d for d in rubric.dimensions if d.name == dimension), None)
                        if dim_def:
                            min_s = dim_def.scale[0]
                            max_s = dim_def.scale[-1]
                            normalized = (raw_score - min_s) / (max_s - min_s) if max_s != min_s else 1.0
                        else:
                            normalized = raw_score / 5.0  # fallback: escala 1-5

                        evaluator_name = f"{ev_conf.name}_{dimension}"
                        results.append((evaluator_name, round(normalized, 4), rationale))

                except Exception as e:
                    self.console.print(
                        f"[red]✗ Judge '{ev_conf.name}' falhou: {e}[/red]"
                    )

            else:
                self.console.print(
                    f"[yellow]⚠ Tipo de evaluador desconhecido: '{ev_conf.type}'. Skipping.[/yellow]"
                )

        return results

    def export_json(self, run_id: str, path: str) -> None:
        from pathlib import Path
        data = {
            "run_id": run_id,
            "cases": self.case_repo.get_by_run(run_id),
            "scores": self.score_repo.get_by_run(run_id),
        }
        Path(path).write_text(json.dumps(data, indent=2), encoding="utf-8")