from __future__ import annotations

import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Literal, overload

from .client import HarnessClient, HarnessConfig, UserQuestionHandler
from .errors import SdkProtocolError, TransportClosedError
from .models import CommandExecution, JsonObject, Notification


@dataclass(slots=True)
class DeepSeekHarnessConfig:
    """Configuration for launching the local DeepSeek Harness SDK runtime.

    The runtime inherits the caller's environment by default, so existing
    DEEPSEEK_API_KEY and DEEPSEEK_BASE_URL settings keep working. Use ``env`` to
    intentionally override or inject variables for a subprocess.
    These values are startup input, not a live Session state mirror; Session
    commands do not rewrite this object.
    """

    provider: str = "deepseek-official"
    model: str = "deepseek-v4-flash"
    reasoning_effort: str | None = None
    max_tokens: int | None = None
    cwd: str | None = None
    runtime_cwd: str | None = None
    dsh_bin: str | None = None
    profile: str = "sdk"
    patches: tuple[str, ...] = ()
    dsh_home: str | None = None
    env: dict[str, str] = field(default_factory=dict)
    initialize_timeout_seconds: float = 30.0
    request_timeout_seconds: float | None = None
    shutdown_timeout_seconds: float | None = 1.0
    base_url: str | None = None
    api_key: str | None = None
    on_user_question: UserQuestionHandler | None = None
    approval_requests: bool = False


@dataclass(slots=True)
class RunResult:
    session_id: str
    final_response: str
    finish_reason: str | None
    events: list[JsonObject]
    notifications: list[Notification]


class _PromptActivity:
    """Collect an ordinary prompt's notifications from its inbox receipt through idle."""

    def __init__(self, session_id: str, message_id: str) -> None:
        self.session_id = session_id
        self.message_id = message_id
        self.received = False
        self.events: list[JsonObject] = []
        self.notifications: list[Notification] = []

    def record(self, notification: Notification, on_notification: Callable[[Notification], None] | None = None) -> bool:
        if not self.received:
            if not _is_inbox_receipt(notification, self.session_id, self.message_id):
                return False
            self.received = True
        self.notifications.append(notification)
        if on_notification is not None:
            on_notification(notification)
        if notification.method == "session.event" and notification.payload.get("sessionId") == self.session_id:
            event = notification.payload.get("event")
            if isinstance(event, dict):
                self.events.append(event)
        return True

    def result(self) -> RunResult:
        return RunResult(
            session_id=self.session_id, final_response=final_response(self.events),
            finish_reason=finish_reason(self.events), events=self.events, notifications=self.notifications,
        )


class DeepSeekHarness:
    """Reusable synchronous SDK for running DeepSeek Harness agent turns.

    The runtime subprocess starts lazily and remains owned by this instance
    across calls to :meth:`run`. Use the instance as a context manager, or call
    :meth:`close` explicitly when finished, so the subprocess is always reaped.
    """

    def __init__(
        self,
        config: DeepSeekHarnessConfig | None = None,
        *,
        _launch_args: tuple[str, ...] | None = None,
        **kwargs: object,
    ) -> None:
        if config is not None and kwargs:
            raise TypeError("pass either DeepSeekHarnessConfig or keyword options, not both")
        self.config = config or DeepSeekHarnessConfig(**kwargs)
        cwd = str(Path(self.config.cwd or Path.cwd()).resolve())
        runtime_cwd = str(Path(self.config.runtime_cwd).resolve()) if self.config.runtime_cwd is not None else cwd
        self._cwd = cwd
        env = dict(self.config.env)
        if self.config.base_url is not None:
            env["DEEPSEEK_BASE_URL"] = self.config.base_url
        if self.config.api_key is not None:
            env["DEEPSEEK_API_KEY"] = self.config.api_key

        self._client = HarnessClient(
            HarnessConfig(
                dsh_bin=self.config.dsh_bin,
                profile=self.config.profile,
                patches=self.config.patches,
                dsh_home=self.config.dsh_home,
                cwd=runtime_cwd,
                env=env,
                initialize_timeout_seconds=self.config.initialize_timeout_seconds,
                request_timeout_seconds=self.config.request_timeout_seconds,
                shutdown_timeout_seconds=self.config.shutdown_timeout_seconds,
                on_user_question=self.config.on_user_question,
            ),
            _launch_args=_launch_args,
        )
        self._initialized = False

    def __enter__(self) -> "DeepSeekHarness":
        self.start()
        return self

    def __exit__(self, _exc_type, _exc, _tb) -> None:
        self.close()

    @property
    def client(self) -> HarnessClient:
        return self._client

    def start(self) -> None:
        if self._initialized:
            return
        self._client.start()
        self._client.initialize(
            cwd=self._cwd,
            provider=self.config.provider,
            model=self.config.model,
            reasoning_effort=self.config.reasoning_effort,
            max_tokens=self.config.max_tokens,
            approval_requests=self.config.approval_requests,
        )
        self._initialized = True

    def close(self) -> None:
        self._client.close()
        self._initialized = False

    def start_session(self, session_id: str | None = None) -> "Session":
        self.start()
        return Session(self, session_id or f"session-{uuid.uuid4().hex}")

    def resume_session(self, session_id: str) -> "Session":
        """Restore persisted history from the configured Harness home without sending input.

        Missing, unreadable, incompatible, or occupied sessions raise instead of creating a new session.
        """
        self.start()
        self._client.resume_session(session_id)
        return Session(self, session_id)

    @overload
    def run(
        self, input: list[JsonObject], *, session_id: str | None = None,
        on_notification: Callable[[Notification], None] | None = None,
    ) -> RunResult: ...

    @overload
    def run(
        self, input: str | list[JsonObject], *, session_id: str | None = None,
        on_notification: Callable[[Notification], None] | None = None,
        delivery: Literal["queue", "steer"] = "queue",
    ) -> RunResult | CommandExecution: ...

    def run(
        self,
        input: str | list[JsonObject],
        *,
        session_id: str | None = None,
        on_notification: Callable[[Notification], None] | None = None,
        delivery: Literal["queue", "steer"] = "queue",
    ) -> RunResult | CommandExecution:
        """Run text or a native command on a fresh or named Session."""
        return self.start_session(session_id).run(input, on_notification=on_notification, delivery=delivery)


class Session:
    def __init__(self, harness: DeepSeekHarness, session_id: str) -> None:
        self.harness = harness
        self.id = session_id

    def get_working_directory(self) -> str:
        """Read the effective directory, recovering a missing directory to the Session origin."""
        self.harness.start()
        return self.harness.client.get_working_directory(self.id)

    def set_working_directory(self, path: str) -> str:
        """Change the effective directory while existing processes retain their own directories."""
        self.harness.start()
        return self.harness.client.set_working_directory(self.id, path)

    @overload
    def run(
        self, input: list[JsonObject], *,
        on_notification: Callable[[Notification], None] | None = None,
    ) -> RunResult: ...

    @overload
    def run(
        self, input: str | list[JsonObject], *,
        on_notification: Callable[[Notification], None] | None = None,
        delivery: Literal["queue", "steer"] = "queue",
    ) -> RunResult | CommandExecution: ...

    def run(
        self,
        input: str | list[JsonObject],
        *,
        on_notification: Callable[[Notification], None] | None = None,
        delivery: Literal["queue", "steer"] = "queue",
    ) -> RunResult | CommandExecution:
        """Run text through native input routing, or queue literal content blocks.

        Text returns the prompt's receipt-to-idle RunResult or a settled
        CommandExecution; command completion does not wait for later model work.
        Content blocks stay literal and accept only queue delivery. Notifications
        retain wire order. Command callbacks replay Session-tree notifications
        from subscription through their own command/done after the response,
        including overlapping work. Handler errors remain command results;
        routing, transport and observer failures raise without rolling back work.
        """
        if delivery not in ("queue", "steer"):
            raise ValueError("run delivery must be queue or steer")
        if not isinstance(input, str) and delivery != "queue":
            raise ValueError("content-block input requires queue delivery")
        self.harness.start()
        with self.harness.client.subscribe_session_notifications(self.id) as subscription:
            if isinstance(input, str):
                response = self.harness.client._request_raw(
                    "session/query", {"sessionId": self.id, "text": input, "delivery": delivery},
                )
                if not isinstance(response, dict):
                    raise SdkProtocolError("session/query response must be a JSON object")
                if response.get("kind") == "command":
                    execution = CommandExecution.model_validate(response.get("execution"))
                    while True:
                        try:
                            notification = subscription._next_queued()
                        except TransportClosedError:
                            break
                        if notification is None:
                            break
                        if on_notification is not None:
                            on_notification(notification)
                        if _is_command_done(notification, self.id, execution.command_id):
                            break
                    return execution
                if response.get("kind") != "prompt" or not isinstance(response.get("messageId"), str) or not response["messageId"]:
                    raise SdkProtocolError("session/query requires a prompt messageId or command execution result")
                message_id = response["messageId"]
            else:
                message_id = self.harness.client.session_prompt(
                    self.id, input, notification_subscription=subscription,
                )
            activity = _PromptActivity(self.id, message_id)
            while True:
                notification = subscription.next()
                if not activity.record(notification, on_notification):
                    continue
                if (
                    notification.method == "session.status"
                    and notification.payload.get("sessionId") == self.id
                    and notification.payload.get("status") == "idle"
                ):
                    break

        return activity.result()


def _is_inbox_receipt(notification: Notification, session_id: str, message_id: str) -> bool:
    if notification.method != "session.event" or notification.payload.get("sessionId") != session_id:
        return False
    event = notification.payload.get("event")
    if not isinstance(event, dict) or event.get("type") != "agent/inbox/spliced":
        return False
    data = event.get("data")
    inserted = data.get("inserted") if isinstance(data, dict) else None
    return isinstance(inserted, list) and any(
        isinstance(message, dict) and message.get("id") == message_id for message in inserted
    )


def _is_command_done(notification: Notification, session_id: str, command_id: str) -> bool:
    if notification.method != "session.event" or notification.payload.get("sessionId") != session_id:
        return False
    event = notification.payload.get("event")
    if not isinstance(event, dict) or event.get("type") != "command/done":
        return False
    data = event.get("data")
    return isinstance(data, dict) and data.get("commandId") == command_id


def normalize_input(input: str | list[JsonObject]) -> list[JsonObject]:
    if isinstance(input, str):
        return [{"type": "text", "text": input}]
    return input


def final_response(events: list[JsonObject]) -> str:
    for event in reversed(events):
        if event.get("type") != "assistant/message":
            continue
        data = event.get("data")
        if not isinstance(data, dict):
            continue
        message = data.get("message")
        content_owner = message if isinstance(message, dict) else data
        content = content_owner.get("content")
        if not isinstance(content, list):
            continue
        parts: list[str] = []
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                parts.append(str(block.get("text") or ""))
        return "".join(parts)
    return ""


def finish_reason(events: list[JsonObject]) -> str | None:
    """Return the last turn-ending kind.

    The input must contain root-session events from one owned run interval.

    Raises:
        SdkProtocolError: The last ``turn/end`` has no string reason kind.
    """
    for event in reversed(events):
        if event.get("type") != "turn/end":
            continue
        data = event.get("data")
        reason = data.get("reason") if isinstance(data, dict) else None
        kind = reason.get("kind") if isinstance(reason, dict) else None
        if not isinstance(kind, str):
            raise SdkProtocolError("turn/end event requires a string data.reason.kind")
        return kind
    return None