from __future__ import annotations

import threading
from datetime import datetime, timedelta, timezone

from trenches_env.benchmark_runner import ScenarioBenchmarkRunner
from trenches_env.rl import DEFAULT_TRAINING_STAGE
from trenches_env.env import FogOfWarDiplomacyEnv
from trenches_env.models import (
    BenchmarkRunRequest,
    BenchmarkRunResponse,
    IngestNewsRequest,
    IngestNewsResponse,
    LiveControlRequest,
    ProviderDiagnosticsResponse,
    ReactionLogEntry,
    ScenarioSummary,
    SessionState,
    SourceMonitorReport,
    StepSessionRequest,
    StepSessionResponse,
)
from trenches_env.source_ingestion import SourceHarvester


class SessionManager:
    def __init__(self, env: FogOfWarDiplomacyEnv | None = None) -> None:
        self.env = env or FogOfWarDiplomacyEnv()
        self._sessions: dict[str, SessionState] = {}
        self._lock = threading.RLock()
        self._background_tick_seconds = 1.0
        self._background_stop = threading.Event()
        self._background_thread: threading.Thread | None = None

    def start_background_runner(self, tick_interval_seconds: float | None = None) -> None:
        with self._lock:
            if tick_interval_seconds is not None:
                self._background_tick_seconds = max(0.05, tick_interval_seconds)
            if self._background_thread is not None and self._background_thread.is_alive():
                return
            self._background_stop.clear()
            self._background_thread = threading.Thread(
                target=self._run_background_loop,
                name="trenches-session-manager",
                daemon=True,
            )
            self._background_thread.start()

    def stop_background_runner(self) -> None:
        self._background_stop.set()
        thread = self._background_thread
        if thread is not None and thread.is_alive():
            thread.join(timeout=max(1.0, self._background_tick_seconds * 2.0))
        self._background_thread = None

    def shutdown(self) -> None:
        self.stop_background_runner()
        self.env.shutdown()

    def create_session(
        self,
        seed: int | None = None,
        training_agent: str = "us",
        training_stage: str = DEFAULT_TRAINING_STAGE,
        max_turns: int | None = None,
        scenario_id: str | None = None,
        replay_id: str | None = None,
        replay_start_index: int | None = None,
    ) -> SessionState:
        with self._lock:
            session = self.env.create_session(
                seed=seed,
                training_agent=training_agent,
                training_stage=training_stage,
                max_turns=max_turns,
                scenario_id=scenario_id,
                replay_id=replay_id,
                replay_start_index=replay_start_index,
            )
            self._sessions[session.session_id] = session
            return session

    def reset_session(
        self,
        session_id: str,
        seed: int | None = None,
        training_agent: str = "us",
        training_stage: str = DEFAULT_TRAINING_STAGE,
        max_turns: int | None = None,
        scenario_id: str | None = None,
        replay_id: str | None = None,
        replay_start_index: int | None = None,
    ) -> SessionState:
        with self._lock:
            self._require_session(session_id)
            session = self.env.reset_session(
                session_id=session_id,
                seed=seed,
                training_agent=training_agent,
                training_stage=training_stage,
                max_turns=max_turns,
                scenario_id=scenario_id,
                replay_id=replay_id,
                replay_start_index=replay_start_index,
            )
            self._sessions[session_id] = session
            return session

    def get_session(self, session_id: str) -> SessionState:
        with self._lock:
            session = self._require_session(session_id)
            if session.live.enabled and session.live.auto_step:
                refreshed = self.env.maybe_auto_step_live_session(session)
            else:
                refreshed = self.env.refresh_session_sources(session)
            self._sessions[session_id] = refreshed
            return refreshed

    def set_live_mode(self, session_id: str, request: LiveControlRequest) -> SessionState:
        with self._lock:
            current = self._require_session(session_id)
            updated = self.env.configure_live_session(current, request)
            self._sessions[session_id] = updated
            return updated

    def step_session(self, session_id: str, request: StepSessionRequest) -> StepSessionResponse:
        with self._lock:
            current = self._require_session(session_id)
            result = self.env.step_session(current, request)
            self._sessions[session_id] = result.session
            return result

    def ingest_news(self, session_id: str, request: IngestNewsRequest) -> IngestNewsResponse:
        with self._lock:
            if not request.signals:
                raise ValueError("At least one external signal is required.")
            current = self._require_session(session_id)
            refreshed = self.env.refresh_session_sources(current)
            actions = self.env.resolve_policy_actions(
                refreshed,
                request.signals,
                agent_ids=request.agent_ids or None,
            )
            result = self.env.step_session(
                refreshed,
                StepSessionRequest(actions=actions, external_signals=request.signals),
            )
            self._sessions[session_id] = result.session
            reaction: ReactionLogEntry | None = result.session.reaction_log[-1] if result.session.reaction_log else None
            return IngestNewsResponse(
                session=result.session,
                oversight=result.oversight,
                reaction=reaction,
                done=result.done,
            )

    def refresh_session_sources(self, session_id: str, force: bool = False) -> SessionState:
        with self._lock:
            current = self._require_session(session_id)
            refreshed = self.env.refresh_session_sources(current, force=force)
            self._sessions[session_id] = refreshed
            return refreshed

    def source_monitor(self, session_id: str) -> SourceMonitorReport:
        with self._lock:
            current = self._require_session(session_id)
            refreshed = self.env.refresh_session_sources(current)
            self._sessions[session_id] = refreshed
            return self.env.source_monitor(refreshed)

    def reaction_log(self, session_id: str) -> list[ReactionLogEntry]:
        with self._lock:
            current = self._require_session(session_id)
            return [entry.model_copy(deep=True) for entry in current.reaction_log]

    def provider_diagnostics(self, session_id: str) -> ProviderDiagnosticsResponse:
        with self._lock:
            current = self._require_session(session_id)
            refreshed = self.env.refresh_session_sources(current)
            self._sessions[session_id] = refreshed
            return self.env.provider_diagnostics(refreshed)

    def list_scenarios(self) -> list[ScenarioSummary]:
        return [
            ScenarioSummary(
                id=scenario.id,
                name=scenario.name,
                description=scenario.description,
                tags=list(scenario.tags),
                benchmark_turns=scenario.benchmark_turns,
                benchmark_enabled=scenario.benchmark_enabled,
            )
            for scenario in self.env.list_scenarios()
        ]

    def run_benchmark(self, request: BenchmarkRunRequest) -> BenchmarkRunResponse:
        runner = ScenarioBenchmarkRunner(
            env_factory=lambda: FogOfWarDiplomacyEnv(source_harvester=SourceHarvester(auto_start=False))
        )
        return runner.run(request)

    def _run_background_loop(self) -> None:
        while not self._background_stop.is_set():
            try:
                self._tick_live_sessions()
            except Exception:
                pass
            self._background_stop.wait(self._background_tick_seconds)

    def _tick_live_sessions(self) -> None:
        now = datetime.now(timezone.utc)
        with self._lock:
            for session_id, session in list(self._sessions.items()):
                if not self._session_needs_live_tick(session, now):
                    continue
                self._sessions[session_id] = self.env.maybe_auto_step_live_session(session)

    @staticmethod
    def _session_needs_live_tick(session: SessionState, now: datetime) -> bool:
        if not session.live.enabled or not session.live.auto_step:
            return False
        if session.live.last_auto_step_at is None:
            return True
        interval = timedelta(milliseconds=max(session.live.poll_interval_ms, 1_000))
        return now - session.live.last_auto_step_at >= interval

    def _require_session(self, session_id: str) -> SessionState:
        session = self._sessions.get(session_id)
        if session is None:
            raise KeyError(session_id)
        return session