"""Repository for evaluation results and turns."""

from datetime import datetime
from typing import Optional

from sqlmodel import Session

from agenteval.models import EvalResult
from agenteval.storage.db import EvalResultDB, TurnDB, get_session


def result_to_db(result: EvalResult) -> EvalResultDB:
    return EvalResultDB(
        id=result.id,
        run_id=result.run_id,
        case_id=result.case_id,
        turn_id=result.turn_id,
        rule_type=result.rule_type,
        passed=result.passed,
        score=result.score,
        reason=result.reason,
    )


def result_from_db(db: EvalResultDB) -> EvalResult:
    return EvalResult(
        id=db.id,
        run_id=db.run_id,
        case_id=db.case_id,
        turn_id=db.turn_id,
        rule_type=db.rule_type,
        passed=db.passed,
        score=db.score,
        reason=db.reason,
    )


class ResultRepository:
    """Repository for evaluation results."""

    def __init__(self, session: Optional[Session] = None):
        self.session = session or get_session()

    def save_turn(self, turn) -> TurnDB:
        db = TurnDB(
            id=turn.id,
            run_id=turn.run_id,
            case_id=turn.case_id,
            round_index=turn.round_index,
            question_msg_id=turn.question_msg_id,
            sent_at=turn.sent_at,
            received_at=turn.received_at,
            latency_ms=turn.latency_ms,
        )
        db.set_sent_message(turn.sent_message)
        db.set_reply(turn.reply)
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        return db

    def update_turn_exchange(
        self,
        turn_id: str,
        *,
        question_msg_id: Optional[str],
        reply: Optional[dict],
        received_at: Optional[datetime],
        latency_ms: Optional[int],
    ) -> Optional[TurnDB]:
        """Attach exchange facts to an already-persisted Turn.

        This is deliberately narrower than updating a whole Turn: the sent
        message, case identity, ordering and send timestamp remain untouched.
        A caller can therefore commit the sent fact before polling and safely
        complete the same ledger row after a reply, timeout or poll failure.
        """
        db = self.session.get(TurnDB, turn_id)
        if db is None:
            return None
        db.question_msg_id = question_msg_id
        db.set_reply(reply)
        db.received_at = received_at
        db.latency_ms = latency_ms
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        return db

    def save_result(self, result: EvalResult) -> EvalResult:
        db = result_to_db(result)
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        return result_from_db(db)