"""GraphMem Python client with retry logic and rate-limit awareness."""

from __future__ import annotations

import random
import time
from typing import Any, Optional
from urllib.parse import quote, urlencode

import httpx

from graphmem.types import (
    BlacklistedTriple,
    Community,
    ConflictLogEntry,
    ConflictResolution,
    ContextOptions,
    ContextResult,
    ContextResultJson,
    ContextResultText,
    DetectCommunitiesResult,
    ExportData,
    FileChunkResult,
    FileSearchResult,
    FileWithContent,
    GraphCommunity,
    GraphData,
    GraphEdge,
    GraphNode,
    HealthResult,
    ImportResult,
    IngestResult,
    MemoryFile,
    PendingFact,
    Project,
    RateLimitInfo,
    RememberResult,
    RetryConfig,
    SearchResult,
    Trace,
    TraceDetail,
    Triple,
    UsageInfo,
    WriteFileResult,
)


class GraphMemError(Exception):
    """Error from the GraphMem API."""

    def __init__(self, message: str, status: int, body: Any = None) -> None:
        super().__init__(message)
        self.status = status
        self.body = body


class GraphMem:
    """Typed client for the GraphMem knowledge graph API.

    Example::

        from graphmem import GraphMem

        mem = GraphMem(api_key="gm_your_key_here")

        mem.remember("Alice works at Acme Corp")
        ctx = mem.get_context("What does Alice do?")
    """

    def __init__(
        self,
        api_key: str,
        *,
        base_url: str = "https://anuramemory.com",
        retry: Optional[RetryConfig] = None,
        timeout: float = 30.0,
    ) -> None:
        if not api_key:
            raise ValueError("graphmem: api_key is required")
        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._retry = retry or RetryConfig()
        self._client = httpx.Client(timeout=timeout)
        self.rate_limit = RateLimitInfo()

    def close(self) -> None:
        """Close the underlying HTTP client."""
        self._client.close()

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

    def __exit__(self, *args: Any) -> None:
        self.close()

    # -------------------------------------------------------------------------
    # Internal fetch helper with retry + rate-limit parsing
    # -------------------------------------------------------------------------

    def _fetch(self, method: str, path: str, body: Any = None) -> Any:
        url = f"{self._base_url}{path}"
        headers: dict[str, str] = {"X-API-KEY": self._api_key}
        if body is not None:
            headers["Content-Type"] = "application/json"

        last_error: Optional[Exception] = None

        for attempt in range(self._retry.max_retries + 1):
            if attempt > 0:
                delay = min(
                    self._retry.base_delay * (2 ** (attempt - 1)),
                    self._retry.max_delay,
                )
                jitter = delay * (0.75 + random.random() * 0.5)
                time.sleep(jitter)

            try:
                response = self._client.request(
                    method,
                    url,
                    headers=headers,
                    json=body if body is not None else None,
                    content=None if body is not None else None,
                )
            except httpx.HTTPError as exc:
                last_error = exc
                if attempt < self._retry.max_retries:
                    continue
                raise

            self._parse_rate_limit_headers(response)

            if (
                response.status_code in self._retry.retry_on
                and attempt < self._retry.max_retries
            ):
                if response.status_code == 429:
                    retry_after = response.headers.get("Retry-After")
                    if retry_after:
                        try:
                            seconds = int(retry_after)
                            time.sleep(seconds)
                            continue
                        except ValueError:
                            pass
                continue

            data = response.json()

            if response.status_code >= 400 or data.get("success") is False:
                raise GraphMemError(
                    data.get("error", f"Request failed: {method} {path}"),
                    response.status_code,
                    data,
                )

            return data

        raise last_error or Exception("Request failed after retries")

    def _parse_rate_limit_headers(self, response: httpx.Response) -> None:
        remaining = response.headers.get("X-RateLimit-Remaining")
        limit = response.headers.get("X-RateLimit-Limit")
        reset = response.headers.get("X-RateLimit-Reset")
        self.rate_limit = RateLimitInfo(
            remaining=int(remaining) if remaining else None,
            limit=int(limit) if limit else None,
            reset=int(reset) if reset else None,
        )

    # -------------------------------------------------------------------------
    # Health
    # -------------------------------------------------------------------------

    def health(self) -> HealthResult:
        """Check server health and connectivity."""
        data = self._fetch("GET", "/api/v1/health")
        return HealthResult(
            status=data["status"],
            version=data["version"],
            timestamp=data["timestamp"],
        )

    # -------------------------------------------------------------------------
    # Core
    # -------------------------------------------------------------------------

    def remember(self, text: str) -> RememberResult:
        """Extract knowledge from text and store in the graph."""
        data = self._fetch("POST", "/api/v1/remember", {"text": text})
        raw_conflicts = data.get("conflicts")
        conflicts: Optional[list[ConflictResolution]] = None
        if raw_conflicts:
            conflicts = [
                ConflictResolution(
                    subject_name=c["subjectName"],
                    predicate=c["predicate"],
                    old_object=c["oldObject"],
                    new_object=c["newObject"],
                    resolution=c["resolution"],
                )
                for c in raw_conflicts
            ]
        return RememberResult(
            extracted_count=data["extractedCount"],
            merged_count=data["mergedCount"],
            message=data["message"],
            conflicts=conflicts,
        )

    def get_context(
        self, query: str, options: Optional[ContextOptions] = None
    ) -> ContextResult:
        """Retrieve graph context for a query."""
        params: dict[str, str] = {"q": query}
        if options and options.format:
            params["format"] = options.format
        if options and options.mode:
            params["mode"] = options.mode

        data = self._fetch("GET", f"/api/v1/context?{urlencode(params)}")

        fmt = data.get("format", "json")
        if fmt == "json":
            return ContextResultJson(
                format="json",
                query=data["query"],
                trace_id=data.get("traceId"),
                entities=data.get("entities", []),
                edges=data.get("edges", []),
                vector_entities=data.get("vectorEntities"),
                communities=data.get("communities"),
            )
        return ContextResultText(
            format=fmt,
            query=data["query"],
            trace_id=data.get("traceId"),
            content=data["content"],
        )

    def search(self, entity: str) -> SearchResult:
        """Search for an entity and its 1-hop neighbors."""
        params = urlencode({"q": entity})
        data = self._fetch("GET", f"/api/v1/search?{params}")
        return SearchResult(entity=data["entity"], edges=data["edges"])

    # -------------------------------------------------------------------------
    # Graph management
    # -------------------------------------------------------------------------

    def get_graph(self) -> GraphData:
        """Get the full graph (nodes, edges, communities)."""
        data = self._fetch("GET", "/api/v1/graph")
        return GraphData(
            nodes=[GraphNode(**n) for n in data.get("nodes", [])],
            edges=[GraphEdge(**e) for e in data.get("edges", [])],
            communities=(
                [GraphCommunity(**c) for c in data["communities"]]
                if data.get("communities")
                else None
            ),
        )

    def ingest_triples(self, triples: list[Triple]) -> IngestResult:
        """Ingest triples directly into the graph."""
        payload = [
            {"subject": t.subject, "predicate": t.predicate, "object": t.object}
            for t in triples
        ]
        data = self._fetch("POST", "/api/v1/graph/triples", {"triples": payload})
        return IngestResult(count=data["count"], message=data["message"])

    def delete_edge(self, id: str, *, blacklist: bool = False) -> None:
        """Delete an edge. Optionally blacklist the triple."""
        params = "?blacklist=true" if blacklist else ""
        self._fetch("DELETE", f"/api/v1/graph/edges/{quote(id, safe='')}{params}")

    def update_edge_weight(
        self,
        id: str,
        *,
        weight: Optional[float] = None,
        increment: Optional[float] = None,
    ) -> None:
        """Update an edge's weight."""
        body: dict[str, Any] = {}
        if weight is not None:
            body["weight"] = weight
        if increment is not None:
            body["increment"] = increment
        self._fetch("PATCH", f"/api/v1/graph/edges/{quote(id, safe='')}", body)

    def delete_node(self, id: str) -> None:
        """Delete a node and all its connected edges."""
        self._fetch("DELETE", f"/api/v1/graph/nodes/{quote(id, safe='')}")

    def export_graph(self) -> ExportData:
        """Export the entire graph as portable JSON."""
        data = self._fetch("GET", "/api/v1/graph/export")
        return ExportData(version=data["version"], triples=data["triples"])

    def import_graph(self, data: ExportData) -> ImportResult:
        """Import a graph export into the current project."""
        res = self._fetch(
            "POST",
            "/api/v1/graph/import",
            {"version": data.version, "triples": data.triples},
        )
        return ImportResult(imported=res["imported"], message=res["message"])

    # -------------------------------------------------------------------------
    # Traces
    # -------------------------------------------------------------------------

    def list_traces(
        self, *, limit: Optional[int] = None, cursor: Optional[str] = None
    ) -> list[Trace]:
        """List query traces with pagination."""
        params: dict[str, str] = {}
        if limit is not None:
            params["limit"] = str(limit)
        if cursor is not None:
            params["cursor"] = cursor
        qs = f"?{urlencode(params)}" if params else ""
        data = self._fetch("GET", f"/api/v1/traces{qs}")
        return [
            Trace(
                id=t["id"],
                query=t["query"],
                format=t["format"],
                mode=t["mode"],
                entity_count=t["entityCount"],
                edge_count=t["edgeCount"],
                created_at=t.get("createdAt", 0),
            )
            for t in data.get("traces", [])
        ]

    def get_trace(self, id: str) -> TraceDetail:
        """Get details for a specific trace."""
        data = self._fetch("GET", f"/api/v1/traces/{quote(id, safe='')}")
        return TraceDetail(
            id=data["id"],
            query=data["query"],
            format=data["format"],
            mode=data["mode"],
            entity_count=data["entityCount"],
            edge_count=data["edgeCount"],
            created_at=data.get("createdAt", 0),
            entities=data.get("entities", []),
            edges=data.get("edges", []),
        )

    # -------------------------------------------------------------------------
    # Blacklist
    # -------------------------------------------------------------------------

    def list_blacklist(
        self, *, limit: Optional[int] = None, cursor: Optional[str] = None
    ) -> list[BlacklistedTriple]:
        """List blacklisted triples with pagination."""
        params: dict[str, str] = {}
        if limit is not None:
            params["limit"] = str(limit)
        if cursor is not None:
            params["cursor"] = cursor
        qs = f"?{urlencode(params)}" if params else ""
        data = self._fetch("GET", f"/api/v1/graph/blacklist{qs}")
        return [
            BlacklistedTriple(
                id=item["id"],
                subject=item["subject"],
                predicate=item["predicate"],
                object=item["object"],
                created_at=item.get("createdAt", 0),
            )
            for item in data.get("items", [])
        ]

    def add_to_blacklist(
        self, subject: str, predicate: str, object: str
    ) -> dict[str, str]:
        """Add a triple to the blacklist."""
        data = self._fetch(
            "POST",
            "/api/v1/graph/blacklist",
            {"subject": subject, "predicate": predicate, "object": object},
        )
        return {"id": data["id"]}

    def remove_from_blacklist(self, id: str) -> None:
        """Remove a triple from the blacklist by ID."""
        self._fetch("DELETE", f"/api/v1/graph/blacklist/{quote(id, safe='')}")

    # -------------------------------------------------------------------------
    # Conflict log
    # -------------------------------------------------------------------------

    def list_conflicts(
        self, *, limit: Optional[int] = None, cursor: Optional[str] = None
    ) -> list[ConflictLogEntry]:
        """List conflict resolution log entries (newest first, paginated)."""
        params: dict[str, str] = {}
        if limit is not None:
            params["limit"] = str(limit)
        if cursor is not None:
            params["cursor"] = cursor
        qs = f"?{urlencode(params)}" if params else ""
        data = self._fetch("GET", f"/api/v1/graph/conflicts{qs}")
        return [
            ConflictLogEntry(
                id=item["id"],
                subject_name=item["subjectName"],
                predicate=item["predicate"],
                old_object=item["oldObject"],
                new_object=item["newObject"],
                resolution=item["resolution"],
                created_at=item.get("createdAt", ""),
            )
            for item in data.get("conflicts", [])
        ]

    # -------------------------------------------------------------------------
    # Pending facts
    # -------------------------------------------------------------------------

    def list_pending(
        self, *, limit: Optional[int] = None, cursor: Optional[str] = None
    ) -> list[PendingFact]:
        """List pending facts with pagination."""
        params: dict[str, str] = {}
        if limit is not None:
            params["limit"] = str(limit)
        if cursor is not None:
            params["cursor"] = cursor
        qs = f"?{urlencode(params)}" if params else ""
        data = self._fetch("GET", f"/api/v1/graph/pending{qs}")
        return [
            PendingFact(
                id=item["id"],
                subject=item["subject"],
                predicate=item["predicate"],
                object=item["object"],
                status=item["status"],
                created_at=item.get("createdAt", 0),
            )
            for item in data.get("items", [])
        ]

    def approve_fact(self, id: str) -> None:
        """Approve a pending fact (merges it into the graph)."""
        self._fetch(
            "POST",
            f"/api/v1/graph/pending/{quote(id, safe='')}",
            {"action": "approve"},
        )

    def reject_fact(self, id: str, *, blacklist: bool = False) -> None:
        """Reject a pending fact (deletes it)."""
        self._fetch(
            "POST",
            f"/api/v1/graph/pending/{quote(id, safe='')}",
            {"action": "reject", "blacklist": blacklist},
        )

    def approve_all(self) -> None:
        """Approve all pending facts at once."""
        self._fetch("POST", "/api/v1/graph/pending", {"action": "approve_all"})

    def reject_all(self) -> None:
        """Reject all pending facts at once."""
        self._fetch("POST", "/api/v1/graph/pending", {"action": "reject_all"})

    # -------------------------------------------------------------------------
    # Projects
    # -------------------------------------------------------------------------

    def list_projects(self) -> list[Project]:
        """List all projects for the authenticated user."""
        data = self._fetch("GET", "/api/v1/projects")
        return [
            Project(
                id=p["id"],
                name=p["name"],
                created_at=p.get("createdAt", 0),
            )
            for p in data.get("projects", [])
        ]

    def create_project(self, name: str) -> Project:
        """Create a new project."""
        data = self._fetch("POST", "/api/v1/projects", {"name": name})
        p = data["project"]
        return Project(id=p["id"], name=p["name"], created_at=p.get("createdAt", 0))

    def delete_project(self, id: str) -> None:
        """Delete a project by ID."""
        self._fetch("DELETE", f"/api/v1/projects/{quote(id, safe='')}")

    def select_project(self, id: str) -> None:
        """Switch the active project context."""
        self._fetch("POST", f"/api/v1/projects/{quote(id, safe='')}/select")

    # -------------------------------------------------------------------------
    # Communities
    # -------------------------------------------------------------------------

    def list_communities(self) -> list[Community]:
        """List all detected communities."""
        data = self._fetch("GET", "/api/v1/communities")
        return [
            Community(
                id=c["id"],
                community_id=c["communityId"],
                name=c["name"],
                summary=c.get("summary", ""),
                member_count=c.get("memberCount", 0),
                created_at=c.get("createdAt"),
                updated_at=c.get("updatedAt"),
            )
            for c in data.get("communities", [])
        ]

    def detect_communities(self) -> DetectCommunitiesResult:
        """Run community detection on the graph."""
        data = self._fetch("POST", "/api/v1/communities")
        return DetectCommunitiesResult(
            community_count=data["communityCount"],
            entity_count=data["entityCount"],
            message=data["message"],
        )

    # -------------------------------------------------------------------------
    # Usage
    # -------------------------------------------------------------------------

    def get_usage(self) -> UsageInfo:
        """Get current usage and tier info."""
        data = self._fetch("GET", "/api/v1/usage")
        u = data["usage"]
        return UsageInfo(
            tier=u["tier"],
            fact_limit=u["factLimit"],
            current_facts=u["currentFacts"],
            nodes_count=u["nodesCount"],
            edges_count=u["edgesCount"],
            file_storage_limit=u.get("fileStorageLimit", 0),
            file_count_limit=u.get("fileCountLimit", 0),
            current_file_storage=u.get("currentFileStorage", 0),
            current_file_count=u.get("currentFileCount", 0),
        )

    # -------------------------------------------------------------------------
    # Files (FilesRag)
    # -------------------------------------------------------------------------

    def _parse_memory_file(self, f: dict[str, Any]) -> MemoryFile:
        return MemoryFile(
            id=f["id"],
            name=f["name"],
            path=f["path"],
            size=f["size"],
            created_at=f.get("createdAt", ""),
            updated_at=f.get("updatedAt", ""),
            chunk_count=f.get("chunkCount"),
        )

    def write_file(
        self, path: str, content: str, *, name: Optional[str] = None
    ) -> WriteFileResult:
        """Create or update a markdown memory file.

        If a file already exists at the given path, its content is replaced.
        The file is automatically chunked and indexed for semantic search.
        """
        body: dict[str, str] = {"path": path, "content": content}
        body["name"] = name or path.rsplit("/", 1)[-1] or "untitled.md"
        data = self._fetch("POST", "/api/v1/files", body)
        return WriteFileResult(
            file=self._parse_memory_file(data["file"]),
            chunk_count=data["chunkCount"],
            created=data.get("created", True),
        )

    def list_files(self) -> list[MemoryFile]:
        """List all memory files in the current project."""
        data = self._fetch("GET", "/api/v1/files")
        return [self._parse_memory_file(f) for f in data.get("files", [])]

    def read_file(self, id: str) -> FileWithContent:
        """Read a file's content by ID."""
        data = self._fetch("GET", f"/api/v1/files/{quote(id, safe='')}")
        f = data["file"]
        return FileWithContent(
            id=f["id"],
            name=f["name"],
            path=f["path"],
            size=f["size"],
            content=f["content"],
            created_at=f.get("createdAt", ""),
            updated_at=f.get("updatedAt", ""),
            chunk_count=f.get("chunkCount"),
        )

    def update_file(
        self, id: str, content: str, *, name: Optional[str] = None
    ) -> WriteFileResult:
        """Update a file's content (re-chunks and re-indexes)."""
        body: dict[str, str] = {"content": content}
        if name is not None:
            body["name"] = name
        data = self._fetch("PUT", f"/api/v1/files/{quote(id, safe='')}", body)
        return WriteFileResult(
            file=self._parse_memory_file(data["file"]),
            chunk_count=data["chunkCount"],
            created=False,
        )

    def delete_file(self, id: str) -> None:
        """Delete a memory file and all its indexed chunks."""
        self._fetch("DELETE", f"/api/v1/files/{quote(id, safe='')}")

    def search_files(
        self, query: str, *, limit: Optional[int] = None, file_id: Optional[str] = None
    ) -> list[FileSearchResult]:
        """Semantic search across memory files. Pass file_id to search within a single file."""
        params: dict[str, str] = {"q": query}
        if limit is not None:
            params["limit"] = str(limit)
        if file_id is not None:
            params["fileId"] = file_id
        qs = urlencode(params)
        data = self._fetch("GET", f"/api/v1/files/search?{qs}")
        results: list[FileSearchResult] = []
        for r in data.get("results", []):
            chunks = [
                FileChunkResult(
                    heading_title=c.get("headingTitle", ""),
                    excerpt=c.get("excerpt", ""),
                    score=c.get("score", 0.0),
                )
                for c in r.get("chunks", [])
            ]
            results.append(FileSearchResult(file=r["file"], chunks=chunks))
        return results
