from __future__ import annotations

import json
from dataclasses import dataclass, field, asdict
from typing import Any, Optional
from datetime import datetime


@dataclass
class AgentResult:
    agent_name: str
    input_data: str
    output_data: str
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())


@dataclass
class MiddlewareContext:
    user_input: str
    session_id: Optional[str] = None
    use_session: bool = True
    processed_input: Optional[str] = None
    current_result: Optional[str] = None
    agent_results: list[AgentResult] = field(default_factory=list)
    is_early_return: bool = False
    early_return_reason: Optional[str] = None
    metadata: dict = field(default_factory=dict)

    def set_processed_input(self, processed: str) -> None:
        self.processed_input = processed

    def get_input_for_agent(self) -> str:
        return self.processed_input if self.processed_input is not None else self.user_input

    def to_json_input(self, agent_name: str) -> str:
        data = {}

        if self.current_result:
            data["previous_agent"] = self.agent_results[-1].agent_name if self.agent_results else "unknown"
            data["previous_output"] = self.current_result
        else:
            data["previous_agent"] = None
            data["previous_output"] = None

        data["user_original_input"] = self.user_input
        data["current_input_for_agent"] = self.get_input_for_agent()

        return json.dumps(data, ensure_ascii=False, indent=2)

    def add_result(self, agent_name: str, input_data: str, output_data: str) -> None:
        self.agent_results.append(AgentResult(
            agent_name=agent_name,
            input_data=input_data,
            output_data=output_data
        ))
        self.current_result = output_data

    def set_early_return(self, reason: str) -> None:
        self.is_early_return = True
        self.early_return_reason = reason

    def to_storage_json(self) -> dict:
        return {
            "user_input": self.user_input,
            "processed_input": self.processed_input,
            "session_id": self.session_id,
            "use_session": self.use_session,
            "current_result": self.current_result,
            "is_early_return": self.is_early_return,
            "early_return_reason": self.early_return_reason,
            "timestamp": datetime.now().isoformat(),
            "agent_results": [asdict(r) for r in self.agent_results]
        }
