"""HTTP client for Armory API."""

import json
from datetime import datetime
from typing import Any, Dict, List, Optional

import httpx

from armory.config import ArmoryConfig, get_config
from armory.models import (
    Run,
    Session,  # Alias for backwards compatibility
    Project,
    Skill,
    FailedAttempt,
    RetrospectiveRequest,
    RetrospectiveResponse,
    AdviseRequest,
    AdviseResponse,
)


class ArmoryClient:
    """Client for interacting with Armory API."""

    def __init__(self, config: Optional[ArmoryConfig] = None):
        """Initialize the Armory client.

        Args:
            config: Optional configuration. Uses global config if not provided.
        """
        self.config = config or get_config()
        self._client = httpx.Client(
            base_url=self.config.server_url,
            headers=self._get_headers(),
            timeout=30.0,
        )

    def _get_headers(self) -> Dict[str, str]:
        """Get request headers."""
        headers = {"Content-Type": "application/json"}
        if self.config.api_key:
            headers["Authorization"] = f"Bearer {self.config.api_key}"
        return headers

    def close(self):
        """Close the HTTP client."""
        self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    # =========================================================================
    # Project API
    # =========================================================================

    def create_project(self, name: str, description: str = "") -> Project:
        """Create a new project."""
        response = self._client.post(
            "/api/projects",
            json={"name": name, "description": description},
        )
        response.raise_for_status()
        return Project(**response.json())

    def get_project(self, project_id: str) -> Project:
        """Get a project by ID."""
        response = self._client.get(f"/api/projects/{project_id}")
        response.raise_for_status()
        return Project(**response.json())

    def list_projects(self) -> List[Project]:
        """List all projects."""
        response = self._client.get("/api/projects")
        response.raise_for_status()
        return [Project(**p) for p in response.json()]

    # =========================================================================
    # Run API (W&B-compatible)
    # =========================================================================

    def create_run(
        self,
        project_id: str,
        name: str,
        goal: str = "",
        model_type: str = "embedding",
        run_type: str = "training",
        base_model: str = "",
        dataset_path: str = "",
        config: Optional[Dict[str, Any]] = None,
        tags: Optional[List[str]] = None,
        group: str = "",
        notes: str = "",
    ) -> Run:
        """Create a new run."""
        response = self._client.post(
            f"/api/projects/{project_id}/runs",
            json={
                "name": name,
                "goal": goal,
                "model_type": model_type,
                "run_type": run_type,
                "base_model": base_model,
                "dataset_path": dataset_path,
                "config": config or {},
                "tags": tags or [],
                "group": group,
                "notes": notes,
            },
        )
        response.raise_for_status()
        return Run(**response.json())

    def get_run(self, run_id: str) -> Run:
        """Get a run by ID."""
        response = self._client.get(f"/api/runs/{run_id}")
        response.raise_for_status()
        return Run(**response.json())

    def list_runs(self, project_id: Optional[str] = None) -> List[Run]:
        """List runs, optionally filtered by project."""
        if project_id:
            response = self._client.get(f"/api/projects/{project_id}/runs")
        else:
            response = self._client.get("/api/runs")
        response.raise_for_status()
        return [Run(**r) for r in response.json()]

    def start_run(self, project_id: str, run_id: str) -> Run:
        """Start a pending run."""
        response = self._client.post(
            f"/api/projects/{project_id}/runs/{run_id}/start"
        )
        response.raise_for_status()
        return Run(**response.json())

    def log_run_metrics(
        self,
        run_id: str,
        metrics: Dict[str, float],
        step: Optional[int] = None,
    ) -> None:
        """Log metrics to a run."""
        payload = {"metrics": metrics}
        if step is not None:
            payload["step"] = step
        response = self._client.post(
            f"/api/runs/{run_id}/log",
            json=payload,
        )
        response.raise_for_status()

    def complete_run(
        self,
        project_id: str,
        run_id: str,
        metrics: Optional[Dict[str, float]] = None,
    ) -> Run:
        """Mark a run as completed with final metrics."""
        response = self._client.post(
            f"/api/projects/{project_id}/runs/{run_id}/complete",
            json={"metrics": metrics or {}},
        )
        response.raise_for_status()
        return Run(**response.json())

    def fail_run(
        self,
        project_id: str,
        run_id: str,
        error_message: str,
    ) -> Run:
        """Mark a run as failed."""
        response = self._client.post(
            f"/api/projects/{project_id}/runs/{run_id}/fail",
            json={"error": error_message},
        )
        response.raise_for_status()
        return Run(**response.json())

    # Backwards compatibility aliases
    def create_session(self, project_id: str, goal: str, **kwargs) -> Run:
        """Create a new session (alias for create_run)."""
        name = kwargs.pop("name", f"session-{datetime.now().strftime('%Y%m%d-%H%M%S')}")
        run_type = kwargs.pop("session_type", "training")
        return self.create_run(project_id=project_id, name=name, goal=goal, run_type=run_type, **kwargs)

    def get_session(self, session_id: str) -> Run:
        """Get a session by ID (alias for get_run)."""
        return self.get_run(session_id)

    def start_session(self, project_id: str, session_id: str) -> Run:
        """Start a pending session (alias for start_run)."""
        return self.start_run(project_id, session_id)

    def complete_session(self, project_id: str, session_id: str, metrics: Dict[str, float]) -> Run:
        """Mark a session as completed (alias for complete_run)."""
        return self.complete_run(project_id, session_id, metrics)

    def fail_session(self, project_id: str, session_id: str, error_message: str) -> Run:
        """Mark a session as failed (alias for fail_run)."""
        return self.fail_run(project_id, session_id, error_message)

    # =========================================================================
    # Agentic API (HuggingFace Blog patterns)
    # =========================================================================

    def advise(
        self,
        goal: str,
        project_id: Optional[str] = None,
        model_type: Optional[str] = None,
        context: Optional[Dict[str, Any]] = None,
    ) -> AdviseResponse:
        """Get advice before starting an experiment.

        "Search before you experiment" pattern from HuggingFace Blog.
        """
        request = AdviseRequest(
            goal=goal,
            project_id=project_id,
            model_type=model_type,
            context=context or {},
        )
        response = self._client.post("/api/advise", json=request.model_dump())
        response.raise_for_status()
        return AdviseResponse(**response.json())

    def retrospective(
        self,
        session_id: str,
        outcome: str,
        metrics: Dict[str, float],
        error_log: Optional[str] = None,
        hyperparameters: Optional[Dict[str, Any]] = None,
        observations: Optional[str] = None,
        auto_create_skill: bool = True,
    ) -> RetrospectiveResponse:
        """Extract learnings from a completed experiment.

        "Auto-capture learnings after experiments" pattern.
        """
        request = RetrospectiveRequest(
            session_id=session_id,
            outcome=outcome,
            metrics=metrics,
            error_log=error_log,
            hyperparameters=hyperparameters,
            observations=observations,
            auto_create_skill=auto_create_skill,
        )
        response = self._client.post("/api/retrospective", json=request.model_dump())
        response.raise_for_status()
        return RetrospectiveResponse(**response.json())

    # =========================================================================
    # Knowledge API
    # =========================================================================

    def list_skills(
        self,
        project_id: Optional[str] = None,
        model_type: Optional[str] = None,
    ) -> List[Skill]:
        """List available skills."""
        params = {}
        if project_id:
            params["project_id"] = project_id
        if model_type:
            params["model_type"] = model_type

        response = self._client.get("/api/skills", params=params)
        response.raise_for_status()
        return [Skill(**s) for s in response.json()]

    def get_skill(self, skill_id: str) -> Skill:
        """Get a skill by ID."""
        response = self._client.get(f"/api/skills/{skill_id}")
        response.raise_for_status()
        return Skill(**response.json())

    def create_skill(self, skill: Skill) -> Skill:
        """Create a new skill."""
        response = self._client.post("/api/skills", json=skill.model_dump())
        response.raise_for_status()
        return Skill(**response.json())

    def create_failed_attempt(self, attempt: FailedAttempt) -> FailedAttempt:
        """Record a failed attempt (anti-pattern)."""
        response = self._client.post("/api/failed-attempts", json=attempt.model_dump())
        response.raise_for_status()
        return FailedAttempt(**response.json())

    def list_failed_attempts(
        self,
        project_id: Optional[str] = None,
        model_type: Optional[str] = None,
    ) -> List[FailedAttempt]:
        """List failed attempts."""
        params = {}
        if project_id:
            params["project_id"] = project_id
        if model_type:
            params["model_type"] = model_type

        response = self._client.get("/api/failed-attempts", params=params)
        response.raise_for_status()
        return [FailedAttempt(**fa) for fa in response.json()]
