from typing import Dict, Any, Optional
import logging
import json
import time
from observability import logger as obs_logger
from observability import components as obs_components

from domain.training.agent_models import PlanOutput, WeeklySummary
from domain.runner.profile import RunnerProfile
from domain.runner.goal import Goal
from llm.base import LLMClient
from pathlib import Path

logger = logging.getLogger(__name__)


class PlanAgent:
    """
    Generates a weekly plan using an LLM based on recent mileage and consistency via the provided LLMClient.
    """

    def __init__(self, llm_client: LLMClient):
        self.llm_client = llm_client
        self.instruction = self._load_instruction("en")

    def _load_instruction(self, language: str = "en") -> str:
        try:
            # Resolve path relative to this file: src/agents/plan/../../prompts/
            base_path = Path(__file__).parent.parent.parent / "prompts"
            filename = f"plan_{language}.txt"
            file_path = base_path / filename

            if not file_path.exists():
                logger.warning(f"Prompt file not found: {file_path}. Falling back to English.")
                file_path = base_path / "plan_en.txt"

            if not file_path.exists():
                logger.error("English prompt file missing!")
                return "You are generating a weekly running plan. Output JSON."

            return file_path.read_text(encoding="utf-8")
        except Exception as e:
            logger.error(f"Error loading prompt for language {language}: {e}")
            return "You are generating a weekly running plan. Output JSON."

    async def run(
        self,
        features: WeeklySummary,
        language: str = "en",
        profile: Optional[RunnerProfile] = None,
        goal: Optional[Goal] = None,
    ) -> Dict[str, Any]:
        with obs_logger.start_span("plan_agent.run", obs_components.AGENT):
            start_time = time.time()
            # Load language-specific instruction
            self.instruction = self._load_instruction(language)

            # Take most recent week from WeeklySummary
            if features.weekly_km:
                # take most recent week by sorting keys (YYYY-WW)
                weeks = sorted(features.weekly_km.items(), reverse=True)
                recent_km = weeks[0][1]
            else:
                recent_km = 0.0

            cons = features.consistency_score

            # Construct Prompt
            prompt = self._construct_prompt(
                recent_km, cons, language=language, profile=profile, goal=goal
            )

            try:
                # Call LLM via Client
                with obs_logger.start_span("plan_agent.llm", obs_components.AGENT):
                    plan_output = await self.llm_client.generate(
                        prompt, instruction=self.instruction, schema=None, name="plan_agent"
                    )

                result = None
                if isinstance(plan_output, PlanOutput):
                    result = {"plan": plan_output.plan}
                elif isinstance(plan_output, dict):
                    if "plan" in plan_output:
                        if isinstance(plan_output["plan"], str):
                            result = {"plan": plan_output["plan"]}
                        else:
                            result = {"plan": json.dumps(plan_output)}
                    else:
                        result = {"plan": json.dumps({"plan": plan_output})}
                elif isinstance(plan_output, str):
                    raw = plan_output.strip()
                    result = {"plan": raw}
                else:
                    logger.error(f"Unexpected response type from LLM: {type(plan_output)}")
                    fail_msg = (
                        "Não foi possível gerar o plano neste momento."
                        if language == "pt-BR"
                        else "Could not generate plan at this time."
                    )
                    result = {"plan": fail_msg}

                return result

            except Exception as e:
                duration_ms = (time.time() - start_time) * 1000
                obs_logger.log_event(
                    "error",
                    f"Failed to generate plan: {e}",
                    event="error",
                    component=obs_components.AGENT,
                    duration_ms=duration_ms,
                )
                logger.error(f"Failed to generate plan with LLM: {e}", exc_info=True)
                fail_msg = (
                    "Não foi possível gerar o plano neste momento devido a um erro."
                    if language == "pt-BR"
                    else "Could not generate plan at this time due to an error."
                )
                return {"plan": fail_msg}

    def _construct_prompt(
        self,
        recent_km: float,
        consistency: int,
        language: str = "en",
        profile: Optional[RunnerProfile] = None,
        goal: Optional[Goal] = None,
    ) -> str:
        is_pt = language == "pt-BR"
        profile_context = ""
        if profile:
            profile_context = "\n**Runner Profile Context:**\n"
            if profile.runner_display_name:
                profile_context += f"- Display Name: {profile.runner_display_name}\n"
            if profile.age:
                profile_context += f"- Age: {profile.age}\n"
            if profile.experience_level:
                profile_context += f"- Experience Level: {profile.experience_level}\n"
            if profile.baseline_weekly_km:
                profile_context += f"- Baseline Weekly KM: {profile.baseline_weekly_km}\n"
            if profile.injury_history_notes:
                profile_context += f"- Injury Notes: {profile.injury_history_notes}\n"

        if goal:
            goal_type_label = goal.type.replace("_", " ").title()
            date_str = goal.target_date.strftime("%Y-%m-%d") if goal.target_date else "N/A"
            profile_context += f"\n**Current Active Goal:**\n"
            profile_context += f"- Type: {goal_type_label}\n"
            profile_context += f"- Target: {goal.target_value} {goal.unit}\n"
            profile_context += f"- Target Date: {date_str}\n"

        if is_pt:
            return f"""
**Perfil do Corredor:**
{profile_context}
- Quilometragem Semanal Recente: {recent_km:.1f} km
- Score de Consistência: {consistency}/100
"""
        return f"""
**Runner Profile:**
{profile_context}
- Recent Weekly Mileage: {recent_km:.1f} km
- Consistency Score: {consistency}/100
"""