from __future__ import annotations

import asyncio
from collections.abc import Sequence
from datetime import UTC, datetime
from uuid import UUID, uuid4

from pydantic import JsonValue

from mantora.casts.models import Cast
from mantora.models.events import ObservedStep, Session, SessionContext
from mantora.models.targets import Target
from mantora.policy.blocker import PendingRequest, PendingStatus
from mantora.store.interface import SessionStore


class MemorySessionStore(SessionStore):
    def __init__(self) -> None:
        self._sessions: dict[UUID, Session] = {}
        self._steps: dict[UUID, list[ObservedStep]] = {}
        self._queues: dict[UUID, asyncio.Queue[ObservedStep]] = {}
        self._casts: dict[UUID, Cast] = {}  # cast_id -> Cast
        self._session_casts: dict[UUID, list[UUID]] = {}  # session_id -> cast_ids
        self._pending: dict[UUID, PendingRequest] = {}  # request_id -> PendingRequest
        self._session_client_ids: dict[UUID, str | None] = {}
        self._client_default_repo_roots: dict[str, str] = {}
        self._targets: dict[UUID, Target] = {}
        self._active_target_id: UUID | None = None

    def create_session(
        self,
        *,
        title: str | None,
        context: SessionContext | None = None,
        client_id: str | None = None,
    ) -> Session:
        session_id = uuid4()
        session = Session(id=session_id, title=title, created_at=datetime.now(UTC), context=context)
        self._sessions[session_id] = session
        self._session_client_ids[session_id] = client_id
        self._steps[session_id] = []
        self._queues[session_id] = asyncio.Queue()
        self._session_casts[session_id] = []
        return session

    def list_sessions(
        self,
        *,
        q: str | None = None,
        tag: str | None = None,
        repo_name: str | None = None,
        branch: str | None = None,
        since: datetime | None = None,
        has_warnings: bool | None = None,
        has_blocks: bool | None = None,
    ) -> Sequence[Session]:
        sessions = list(self._sessions.values())

        if since is not None:
            sessions = [s for s in sessions if s.created_at >= since]

        def matches(s: Session) -> bool:
            ctx = s.context
            if tag is not None and ((ctx is None) or (ctx.tag != tag)):
                return False
            if repo_name is not None and ((ctx is None) or (ctx.repo_name != repo_name)):
                return False
            if branch is not None and ((ctx is None) or (ctx.branch != branch)):
                return False

            if has_warnings is not None:
                has_any = any(bool(step.warnings) for step in self._steps.get(s.id, []))
                if has_warnings != has_any:
                    return False

            if has_blocks is not None:
                has_any = any(step.kind == "blocker" for step in self._steps.get(s.id, []))
                if has_blocks != has_any:
                    return False

            if q is not None and q.strip():
                needle = q.strip().lower()
                fields = [
                    s.title or "",
                    ctx.repo_name if ctx and ctx.repo_name else "",
                    ctx.branch if ctx and ctx.branch else "",
                    ctx.tag if ctx and ctx.tag else "",
                ]
                if not any(needle in v.lower() for v in fields):
                    return False

            return True

        filtered = [s for s in sessions if matches(s)]
        return sorted(filtered, key=lambda s: s.created_at, reverse=True)

    def get_session(self, session_id: UUID) -> Session | None:
        return self._sessions.get(session_id)

    def update_session_tag(self, session_id: UUID, *, tag: str | None) -> Session | None:
        session = self._sessions.get(session_id)
        if session is None:
            return None

        context = session.context
        if context is None:
            context = SessionContext(tag=tag) if tag is not None else None
        else:
            context = SessionContext(**context.model_dump(exclude={"tag"}), tag=tag)

        updated = Session(**session.model_dump(exclude={"context"}), context=context)
        self._sessions[session_id] = updated
        return updated

    def update_session_context(
        self, session_id: UUID, *, context: SessionContext | None
    ) -> Session | None:
        session = self._sessions.get(session_id)
        if session is None:
            return None
        updated = Session(**session.model_dump(exclude={"context"}), context=context)
        self._sessions[session_id] = updated
        return updated

    def get_session_client_id(self, session_id: UUID) -> str | None:
        client_id = self._session_client_ids.get(session_id)
        return client_id if isinstance(client_id, str) and client_id.strip() else None

    def get_client_default_repo_root(self, client_id: str) -> str | None:
        repo_root = self._client_default_repo_roots.get(client_id)
        return repo_root if isinstance(repo_root, str) and repo_root.strip() else None

    def set_client_default_repo_root(self, client_id: str, *, repo_root: str | None) -> None:
        cleaned = repo_root.strip() if isinstance(repo_root, str) else ""
        if not cleaned:
            self._client_default_repo_roots.pop(client_id, None)
            return
        self._client_default_repo_roots[client_id] = cleaned

    def session_exists(self, session_id: UUID) -> bool:
        """Check if a session exists without fetching full session data."""
        return session_id in self._sessions

    def get_last_active_at(self, session_id: UUID) -> datetime | None:
        """Get the timestamp of the last activity in a session."""
        if session_id not in self._sessions:
            return None

        steps = self._steps.get(session_id, [])
        if not steps:
            return None

        # Return the most recent step timestamp
        return max(step.created_at for step in steps)

    def delete_session(self, session_id: UUID) -> bool:
        if session_id in self._sessions:
            del self._sessions[session_id]
            self._session_client_ids.pop(session_id, None)
            if session_id in self._steps:
                del self._steps[session_id]
            if session_id in self._queues:
                del self._queues[session_id]

            # Clean up casts
            cast_ids = self._session_casts.get(session_id, [])
            for cid in cast_ids:
                if cid in self._casts:
                    del self._casts[cid]
            if session_id in self._session_casts:
                del self._session_casts[session_id]

            # Clean up pending requests
            pending_ids = [p.id for p in self._pending.values() if p.session_id == session_id]
            for pid in pending_ids:
                del self._pending[pid]

            return True
        return False

    def add_step(self, step: ObservedStep) -> None:
        if step.session_id not in self._sessions:
            raise KeyError(step.session_id)

        self._steps[step.session_id].append(step)
        self._queues[step.session_id].put_nowait(step)

    def update_step(
        self,
        step_id: UUID,
        *,
        summary: str | None = None,
        status: str | None = None,
        args: dict[str, JsonValue] | None = None,
    ) -> bool:
        """Update an existing step's fields."""
        for session_id, steps in self._steps.items():
            for i, step in enumerate(steps):
                if step.id == step_id:
                    # Create updated step
                    step_args = step.args if isinstance(step.args, dict) else {}
                    updated_args = {**step_args, **(args or {})}
                    updated_status = status if status is not None else step.status
                    # Validate status is a valid literal
                    if updated_status not in ("ok", "error"):
                        return False

                    updated_decision = step.decision
                    decision_value = updated_args.get("decision")
                    if isinstance(decision_value, str):
                        updated_decision = decision_value  # type: ignore[assignment]

                    updated_step = ObservedStep(
                        **step.model_dump(exclude={"summary", "status", "args", "decision"}),
                        summary=summary if summary is not None else step.summary,
                        status=updated_status,  # type: ignore
                        args=updated_args,
                        decision=updated_decision,
                    )
                    self._steps[session_id][i] = updated_step
                    # Notify via queue
                    self._queues[session_id].put_nowait(updated_step)
                    return True
        return False

    def list_steps(self, session_id: UUID) -> Sequence[ObservedStep]:
        return list(self._steps.get(session_id, []))

    def get_step_queue(self, session_id: UUID) -> asyncio.Queue[ObservedStep] | None:
        return self._queues.get(session_id)

    def add_cast(self, cast: Cast) -> None:
        if cast.session_id not in self._sessions:
            raise KeyError(cast.session_id)

        self._casts[cast.id] = cast
        self._session_casts[cast.session_id].append(cast.id)

    def list_casts(self, session_id: UUID) -> Sequence[Cast]:
        cast_ids = self._session_casts.get(session_id, [])
        return [self._casts[cid] for cid in cast_ids if cid in self._casts]

    def get_cast(self, cast_id: UUID) -> Cast | None:
        return self._casts.get(cast_id)

    def create_pending_request(
        self,
        *,
        request_id: UUID | None = None,
        session_id: UUID,
        tool_name: str,
        arguments: JsonValue | None,
        classification: str | None,
        risk_level: str | None,
        reason: str | None,
        blocker_step_id: UUID | None,
    ) -> PendingRequest:
        from datetime import UTC, datetime
        from uuid import uuid4

        if session_id not in self._sessions:
            raise KeyError(session_id)

        req_id = request_id or uuid4()
        req = PendingRequest(
            id=req_id,
            session_id=session_id,
            created_at=datetime.now(UTC),
            tool_name=tool_name,
            arguments=arguments,
            classification=classification,
            risk_level=risk_level,
            reason=reason,
            blocker_step_id=blocker_step_id,
            status=PendingStatus.pending,
            decided_at=None,
        )
        self._pending[req.id] = req
        return req

    def get_pending_request(self, request_id: UUID) -> PendingRequest | None:
        return self._pending.get(request_id)

    def list_pending_requests(
        self, session_id: UUID, *, status: PendingStatus | None = None
    ) -> Sequence[PendingRequest]:
        items = [p for p in self._pending.values() if p.session_id == session_id]
        if status is not None:
            items = [p for p in items if p.status == status]
        return sorted(items, key=lambda p: p.created_at)

    def decide_pending_request(
        self, request_id: UUID, *, status: PendingStatus
    ) -> PendingRequest | None:
        from datetime import UTC, datetime

        existing = self._pending.get(request_id)
        if existing is None:
            return None
        decided = PendingRequest(
            **existing.model_dump(exclude={"status", "decided_at"}),
            status=status,
            decided_at=datetime.now(UTC),
        )
        self._pending[request_id] = decided

        return decided

    def create_target(
        self,
        *,
        name: str,
        type: str,
        command: list[str] | None = None,
        env: dict[str, str] | None = None,
    ) -> Target:
        target = Target(
            id=uuid4(),
            name=name,
            type=type,
            command=command or [],
            env=env or {},
            created_at=datetime.now(UTC),
            updated_at=datetime.now(UTC),
        )
        self._targets[target.id] = target
        return target

    def list_targets(self) -> Sequence[Target]:
        return list(self._targets.values())

    def get_target(self, target_id: UUID) -> Target | None:
        return self._targets.get(target_id)

    def get_active_target(self) -> Target | None:
        if self._active_target_id:
            return self._targets.get(self._active_target_id)
        return None

    def update_target(
        self,
        target_id: UUID,
        *,
        name: str | None = None,
        type: str | None = None,
        command: list[str] | None = None,
        env: dict[str, str] | None = None,
    ) -> Target | None:
        existing = self._targets.get(target_id)
        if not existing:
            return None

        updated = Target(
            id=existing.id,
            name=name if name is not None else existing.name,
            type=type if type is not None else existing.type,
            command=command if command is not None else existing.command,
            env=env if env is not None else existing.env,
            created_at=existing.created_at,
            updated_at=datetime.now(UTC),
        )
        self._targets[target_id] = updated
        return updated

    def set_active_target(self, target_id: UUID) -> Target | None:
        if target_id not in self._targets:
            return None
        self._active_target_id = target_id
        return self._targets[target_id]

    def delete_target(self, target_id: UUID) -> bool:
        if target_id in self._targets:
            del self._targets[target_id]
            if self._active_target_id == target_id:
                self._active_target_id = None
            return True
        return False
