from __future__ import annotations

from dataclasses import dataclass, field
import json
from typing import Any, Dict, List, Optional
import time
import uuid

from .dag_parser import DAGParser, ParseError, PlanGraph
from .llm_client import LLMClient, LLMClientError
from .symbolic_verifier import SymbolicVerifier, VerificationResult


@dataclass
class RetryConfig:
    max_attempts: int = 3
    backoff_ms: int = 100
    temperature_escalation: bool = True
    temperatures: List[float] = field(default_factory=lambda: [0.3, 0.5, 0.7])


@dataclass
class PlanContext:
    robot_id: str
    current_location: str
    battery_level: int
    available_actions: List[str]
    venue_id: Optional[str] = None
    mission_type: Optional[str] = None
    additional_constraints: Dict[str, Any] = field(default_factory=dict)


@dataclass
class PlanRequest:
    intent: str
    context: PlanContext
    request_id: str
    timeout_ms: int = 30000


@dataclass
class PlanResult:
    success: bool
    plan: Optional[Dict[str, Any]]
    is_safe_noop: bool
    attempts: int
    total_latency_ms: int
    trace_id: str
    failure_reason: Optional[str] = None
    last_violation: Optional[VerificationResult] = None


class PlannerService:
    def __init__(
        self,
        llm_client: Optional[LLMClient] = None,
        dag_parser: Optional[DAGParser] = None,
        verifier: Optional[SymbolicVerifier] = None,
        retry_config: Optional[RetryConfig] = None,
    ) -> None:
        self.llm_client = llm_client or LLMClient()
        self.dag_parser = dag_parser or DAGParser()
        self.verifier = verifier or SymbolicVerifier()
        self.retry_config = retry_config or RetryConfig()

    def plan(self, request: PlanRequest) -> PlanResult:
        start_time = time.monotonic()
        trace_id = request.request_id or uuid.uuid4().hex
        last_feedback: Optional[str] = None
        last_violation: Optional[VerificationResult] = None

        for attempt in range(self.retry_config.max_attempts):
            temperature = self._temperature_for_attempt(attempt)
            try:
                raw_json = self.llm_client.generate(
                    intent=request.intent,
                    context=request.context,
                    feedback=last_feedback,
                    temperature=temperature,
                )
            except LLMClientError as exc:
                return PlanResult(
                    success=False,
                    plan=None,
                    is_safe_noop=False,
                    attempts=attempt + 1,
                    total_latency_ms=self._elapsed_ms(start_time),
                    trace_id=trace_id,
                    failure_reason=f"llm_unavailable: {exc}",
                    last_violation=None,
                )

            parse_result = self.dag_parser.parse(raw_json, expected_intent=request.intent)
            if not parse_result.ok:
                last_feedback = self._format_parse_error(parse_result.error)
                self._sleep_backoff()
                continue

            verification = self.verifier.verify(parse_result.plan_graph, request.context)
            last_violation = verification
            if verification.passed:
                return PlanResult(
                    success=True,
                    plan=parse_result.plan_graph.plan,
                    is_safe_noop=False,
                    attempts=attempt + 1,
                    total_latency_ms=self._elapsed_ms(start_time),
                    trace_id=trace_id,
                    failure_reason=None,
                    last_violation=None,
                )

            last_feedback = self.verifier.format_feedback(verification)
            self._sleep_backoff()

        safe_noop = self._build_safe_noop_plan(request.intent, last_feedback, attempts=self.retry_config.max_attempts)
        return PlanResult(
            success=False,
            plan=safe_noop,
            is_safe_noop=True,
            attempts=self.retry_config.max_attempts,
            total_latency_ms=self._elapsed_ms(start_time),
            trace_id=trace_id,
            failure_reason="max_retries_exceeded",
            last_violation=last_violation,
        )

    def refine_plan(self, intent: str, context: PlanContext, feedback: Dict[str, Any]) -> PlanGraph:
        raw_json = self.llm_client.generate(
            intent=intent,
            context=context,
            feedback=json.dumps(feedback),
            temperature=self.retry_config.temperatures[-1],
        )
        parse_result = self.dag_parser.parse(raw_json, expected_intent=intent)
        if not parse_result.ok:
            raise ValueError(self._format_parse_error(parse_result.error))
        verification = self.verifier.verify(parse_result.plan_graph, context)
        if not verification.passed:
            raise ValueError(self.verifier.format_feedback(verification))
        return parse_result.plan_graph

    def _format_parse_error(self, error: Optional[ParseError]) -> str:
        if not error:
            return "Parsing failed with unknown error."
        parts = [error.message]
        if error.details:
            parts.append(f"DETAILS: {error.details}")
        if error.hint:
            parts.append(f"FIX: {error.hint}")
        return "\n".join(parts)

    def _build_safe_noop_plan(self, intent: str, last_feedback: Optional[str], attempts: int) -> Dict[str, Any]:
        plan_id = f"plan_safe_noop_{uuid.uuid4().hex[:8]}"
        message = (
            "I could not generate a safe plan for your request. "
            "Please try rephrasing or contact the operator."
        )
        return {
            "intent": intent,
            "response_text": message,
            "workflow_topology": "SEQUENTIAL",
            "plan_id": plan_id,
            "steps": [
                {
                    "step_id": "s1",
                    "type": "EXECUTION",
                    "agent_role": "robot",
                    "action": "SPEAK",
                    "params": {"message": message},
                },
                {
                    "step_id": "s2",
                    "type": "EXECUTION",
                    "agent_role": "robot",
                    "action": "ALERT_OPERATOR",
                    "dependencies": ["s1"],
                    "params": {
                        "reason": "plan_generation_failed",
                        "original_intent": intent,
                        "last_violation": last_feedback or "",
                        "attempts": attempts,
                    },
                },
            ],
        }

    def _temperature_for_attempt(self, attempt: int) -> float:
        temps = self.retry_config.temperatures
        if not temps:
            return 0.3
        if not self.retry_config.temperature_escalation:
            return temps[0]
        return temps[min(attempt, len(temps) - 1)]

    def _sleep_backoff(self) -> None:
        if self.retry_config.backoff_ms <= 0:
            return
        time.sleep(self.retry_config.backoff_ms / 1000.0)

    @staticmethod
    def _elapsed_ms(start_time: float) -> int:
        return int((time.monotonic() - start_time) * 1000)