from __future__ import annotations

import json
import time
from typing import Iterator

import httpx

from .errors import SharedAIPlatformError
from .types import PollOptions, RequestContext, SSEEvent, SyncRecord


class SharedAIPlatformClient:
    def __init__(
        self,
        *,
        base_url: str = "http://127.0.0.1:8011",
        api_key: str,
        app_id: str,
        timeout_seconds: float = 30.0,
        client: httpx.Client | None = None,
        api_prefix: str = "/api/v0.2",
    ):
        self.base_url = base_url.rstrip("/")
        self.api_prefix = self._normalize_api_prefix(api_prefix)
        self.api_key = api_key
        self.app_id = app_id
        self._owns_client = client is None
        self._client = client or httpx.Client(base_url=self.base_url, timeout=timeout_seconds)

    def close(self):
        if self._owns_client:
            self._client.close()

    def __enter__(self) -> "SharedAIPlatformClient":
        return self

    def __exit__(self, exc_type, exc, tb):
        self.close()

    def list_tasks(self, context: RequestContext) -> list[dict]:
        return self._request_json("GET", self._path("/tasks"), context=context)["data"]

    def run_task(
        self,
        context: RequestContext,
        task_key: str,
        *,
        inputs: dict,
        provider: str | None = None,
        metadata: dict | None = None,
    ) -> dict:
        return self._request_json(
            "POST",
            self._path(f"/tasks/{task_key}/run"),
            context=context,
            json_body={"inputs": inputs, "provider": provider, "metadata": metadata or {}},
        )["data"]

    def enqueue_task(
        self,
        context: RequestContext,
        task_key: str,
        *,
        inputs: dict,
        provider: str | None = None,
        metadata: dict | None = None,
    ) -> dict:
        return self._request_json(
            "POST",
            self._path(f"/tasks/{task_key}/enqueue"),
            context=context,
            json_body={"inputs": inputs, "provider": provider, "metadata": metadata or {}},
        )["data"]

    def get_task_run(self, context: RequestContext, run_id: str) -> dict:
        return self._request_json("GET", self._path(f"/task-runs/{run_id}"), context=context)["data"]

    def list_jobs(self, context: RequestContext, *, status_value: str | None = None, limit: int = 50) -> list[dict]:
        params = {"limit": limit}
        if status_value:
            params["status_value"] = status_value
        return self._request_json("GET", self._path("/jobs"), context=context, params=params)["data"]

    def get_job(self, context: RequestContext, job_id: str) -> dict:
        return self._request_json("GET", self._path(f"/jobs/{job_id}"), context=context)["data"]

    def run_next_job(self, context: RequestContext) -> dict | None:
        return self._request_json("POST", self._path("/jobs/run-next"), context=context)["data"]["job"]

    def poll_job_until_terminal(self, context: RequestContext, job_id: str, options: PollOptions | None = None) -> dict:
        opts = options or PollOptions()
        deadline = time.time() + opts.timeout_seconds
        while True:
            job = self.get_job(context, job_id)
            if job["status"] in opts.terminal_statuses:
                return job
            if time.time() >= deadline:
                raise TimeoutError(f"job {job_id} did not reach a terminal state within {opts.timeout_seconds} seconds")
            time.sleep(opts.interval_seconds)

    def list_conversations(self, context: RequestContext, *, page: int = 1, page_size: int = 20) -> dict:
        return self._request_json(
            "GET",
            self._path("/conversations"),
            context=context,
            params={"page": page, "page_size": page_size},
        )["data"]

    def get_conversation_detail(self, context: RequestContext, conversation_id: str) -> dict:
        return self._request_json("GET", self._path(f"/conversations/{conversation_id}"), context=context)["data"]

    def list_knowledge_bases(self, context: RequestContext, *, page: int = 1, page_size: int = 20) -> dict:
        return self._request_json(
            "GET",
            self._path("/knowledge-bases"),
            context=context,
            params={"page": page, "page_size": page_size},
        )["data"]

    def list_documents(self, context: RequestContext, knowledge_base_id: str, *, page: int = 1, page_size: int = 20) -> dict:
        return self._request_json(
            "GET",
            self._path(f"/knowledge-bases/{knowledge_base_id}/documents"),
            context=context,
            params={"page": page, "page_size": page_size},
        )["data"]

    def create_knowledge_base(
        self,
        context: RequestContext,
        *,
        name: str,
        description: str | None = None,
        retrieval_mode: str = "hybrid",
    ) -> dict:
        return self._request_json(
            "POST",
            self._path("/knowledge-bases"),
            context=context,
            json_body={"name": name, "description": description, "retrieval_mode": retrieval_mode},
        )["data"]

    def list_knowledge_base_versions(self, context: RequestContext, knowledge_base_id: str) -> list[dict]:
        return self._request_json(
            "GET",
            self._path(f"/knowledge-bases/{knowledge_base_id}/versions"),
            context=context,
        )["data"]

    def publish_knowledge_base(self, context: RequestContext, knowledge_base_id: str) -> dict:
        return self._request_json(
            "POST",
            self._path(f"/knowledge-bases/{knowledge_base_id}/publish"),
            context=context,
        )["data"]

    def unpublish_knowledge_base(self, context: RequestContext, knowledge_base_id: str) -> dict:
        return self._request_json(
            "POST",
            self._path(f"/knowledge-bases/{knowledge_base_id}/unpublish"),
            context=context,
        )["data"]

    def rollback_knowledge_base(
        self,
        context: RequestContext,
        knowledge_base_id: str,
        *,
        target_version_label: str,
        change_summary: str | None = None,
    ) -> dict:
        return self._request_json(
            "POST",
            self._path(f"/knowledge-bases/{knowledge_base_id}/rollback"),
            context=context,
            json_body={"target_version_label": target_version_label, "change_summary": change_summary},
        )["data"]

    def upload_document(
        self,
        context: RequestContext,
        knowledge_base_id: str,
        *,
        file_name: str,
        mime_type: str | None,
        content: bytes,
    ) -> dict:
        response = self._client.post(
            self._path(f"/knowledge-bases/{knowledge_base_id}/documents"),
            headers=self._headers(context),
            files={"file": (file_name, content, mime_type or "application/octet-stream")},
        )
        self._raise_for_status(response)
        payload = response.json()
        if payload.get("code") not in (0, None):
            raise SharedAIPlatformError(
                status_code=response.status_code,
                message=str(payload.get("message") or "request failed"),
                code=payload.get("code"),
                data=payload.get("data"),
            )
        return payload["data"]

    def retry_document(
        self,
        context: RequestContext,
        knowledge_base_id: str,
        document_id: str,
        *,
        file_name: str | None = None,
        mime_type: str | None = None,
        content: bytes | None = None,
    ) -> dict:
        request_kwargs = {"headers": self._headers(context)}
        if content is not None:
            request_kwargs["files"] = {"file": (file_name or "document.txt", content, mime_type or "application/octet-stream")}
        response = self._client.post(
            self._path(f"/knowledge-bases/{knowledge_base_id}/documents/{document_id}/retry"),
            **request_kwargs,
        )
        self._raise_for_status(response)
        payload = response.json()
        if payload.get("code") not in (0, None):
            raise SharedAIPlatformError(
                status_code=response.status_code,
                message=str(payload.get("message") or "request failed"),
                code=payload.get("code"),
                data=payload.get("data"),
            )
        return payload["data"]

    def sync_content(
        self,
        context: RequestContext,
        knowledge_base_id: str,
        *,
        records: list[SyncRecord | dict],
        schema_version: str = "content-sync.v1",
        sync_mode: str = "incremental",
        batch_id: str | None = None,
        cursor_token: str | None = None,
    ) -> dict:
        normalized_records = [record.to_dict() if isinstance(record, SyncRecord) else record for record in records]
        return self._request_json(
            "POST",
            self._path(f"/knowledge-bases/{knowledge_base_id}/sync-content"),
            context=context,
            json_body={
                "schema_version": schema_version,
                "sync_mode": sync_mode,
                "batch_id": batch_id,
                "cursor_token": cursor_token,
                "records": normalized_records,
            },
        )["data"]

    def list_sync_jobs(self, context: RequestContext, knowledge_base_id: str) -> dict:
        return self._request_json("GET", self._path(f"/knowledge-bases/{knowledge_base_id}/sync-jobs"), context=context)["data"]

    def run_legal_review(
        self,
        context: RequestContext,
        *,
        title: str,
        text: str,
        contract_type: str | None = None,
        review_as: str = "party_a",
        provider: str | None = None,
    ) -> dict:
        return self._request_json(
            "POST",
            self._path("/domain-packs/legal-review/reviews/run"),
            context=context,
            json_body={
                "title": title,
                "contract_type": contract_type,
                "review_as": review_as,
                "text": text,
                "provider": provider,
            },
        )["data"]

    def stream_conversation(
        self,
        context: RequestContext,
        *,
        message: str,
        chat_context: dict | None = None,
        provider: str | None = None,
        template_key: str | None = None,
        conversation_id: str | None = None,
    ) -> Iterator[SSEEvent]:
        headers = self._headers(context)
        with self._client.stream(
            "POST",
            self._path("/conversations/stream"),
            headers=headers,
            json={
                "conversation_id": conversation_id,
                "message": message,
                "chat_context": chat_context or {},
                "provider": provider,
                "template_key": template_key,
            },
        ) as response:
            self._raise_for_status(response)
            yield from self._iter_sse_events(response)

    def _iter_sse_events(self, response: httpx.Response) -> Iterator[SSEEvent]:
        data_lines: list[str] = []

        for line in response.iter_lines():
            if isinstance(line, bytes):
                line = line.decode("utf-8")

            if line == "":
                payload = self._parse_sse_payload(data_lines)
                data_lines = []
                if payload is None:
                    continue
                yield SSEEvent(event_type=str(payload.get("type") or "unknown"), payload=payload)
                continue

            if line.startswith("data:"):
                data_part = line.removeprefix("data:").lstrip()
                if data_part:
                    data_lines.append(data_part)
                continue

            if data_lines or line.startswith("{") or line.startswith("["):
                data_lines.append(line)

        payload = self._parse_sse_payload(data_lines)
        if payload is not None:
            yield SSEEvent(event_type=str(payload.get("type") or "unknown"), payload=payload)

    @staticmethod
    def _parse_sse_payload(data_lines: list[str]) -> dict | None:
        if not data_lines:
            return None
        payload_text = "\n".join(part for part in data_lines if part).strip()
        if not payload_text:
            return None
        return json.loads(payload_text)

    def _request_json(
        self,
        method: str,
        path: str,
        *,
        context: RequestContext | None = None,
        json_body: dict | None = None,
        params: dict | None = None,
    ) -> dict:
        response = self._client.request(method, path, headers=self._headers(context), json=json_body, params=params)
        self._raise_for_status(response)
        payload = response.json()
        if payload.get("code") not in (0, None):
            raise SharedAIPlatformError(
                status_code=response.status_code,
                message=str(payload.get("message") or "request failed"),
                code=payload.get("code"),
                data=payload.get("data"),
            )
        return payload

    def _headers(self, context: RequestContext | None = None) -> dict[str, str]:
        headers: dict[str, str] = {}
        if self.api_key.strip():
            headers["Authorization"] = f"Bearer {self.api_key}"
        if context is not None:
            headers.update(
                {
                    "X-Shared-AI-App-Id": self.app_id,
                    "X-Shared-AI-Tenant-Id": context.tenant_id,
                    "X-Shared-AI-User-Id": context.user_id,
                    "X-Shared-AI-User-Role": context.user_role,
                }
            )
        return headers

    def _raise_for_status(self, response: httpx.Response):
        try:
            response.raise_for_status()
        except httpx.HTTPStatusError as exc:
            if not response.is_closed:
                response.read()
            message = response.text
            try:
                payload = response.json()
                message = str(payload.get("message") or message)
                code = payload.get("code")
                data = payload.get("data")
            except Exception:
                code = None
                data = None
            raise SharedAIPlatformError(response.status_code, message, code, data) from exc

    @staticmethod
    def _normalize_api_prefix(api_prefix: str) -> str:
        normalized = (api_prefix or "/api/v0.2").strip().rstrip("/")
        if not normalized.startswith("/"):
            normalized = f"/{normalized}"
        return normalized

    def _path(self, suffix: str) -> str:
        normalized_suffix = suffix if suffix.startswith("/") else f"/{suffix}"
        return f"{self.api_prefix}{normalized_suffix}"