from __future__ import annotations

import json
import math
import sqlite3
import threading
from typing import Any

from ai_orchestrator.config.settings import Settings
from ai_orchestrator.schemas.task_state import TaskState, utc_now
from ai_orchestrator.schemas.workflow_models import MemoryHit
from ai_orchestrator.services.model_router import ModelRouter


class MemoryService:
    def __init__(self, settings: Settings, model_router: ModelRouter) -> None:
        self.settings = settings
        self.model_router = model_router
        self.database_path = settings.memory_database_path
        self.database_path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()
        self._initialize()

    def _connect(self) -> sqlite3.Connection:
        connection = sqlite3.connect(self.database_path)
        connection.row_factory = sqlite3.Row
        return connection

    def _initialize(self) -> None:
        with self._connect() as connection:
            connection.executescript(
                """
                CREATE TABLE IF NOT EXISTS task_memory (
                    task_id TEXT PRIMARY KEY,
                    goal TEXT NOT NULL,
                    repo TEXT NOT NULL,
                    solution_summary TEXT NOT NULL,
                    success INTEGER NOT NULL,
                    task_mode TEXT,
                    validation_failures TEXT,
                    issue_patterns TEXT,
                    focus_files TEXT,
                    files_modified TEXT NOT NULL,
                    timestamp TEXT NOT NULL,
                    embedding TEXT NOT NULL
                );
                CREATE TABLE IF NOT EXISTS repo_knowledge (
                    repo_name TEXT PRIMARY KEY,
                    payload TEXT NOT NULL,
                    embedding TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                );
                CREATE TABLE IF NOT EXISTS pattern_memory (
                    error_signature TEXT PRIMARY KEY,
                    pattern_text TEXT NOT NULL,
                    occurrence_count INTEGER NOT NULL,
                    embedding TEXT NOT NULL,
                    last_seen TEXT NOT NULL
                );
                """
            )
            self._ensure_column(connection, "task_memory", "task_mode", "TEXT")
            self._ensure_column(connection, "task_memory", "validation_failures", "TEXT")
            self._ensure_column(connection, "task_memory", "issue_patterns", "TEXT")
            self._ensure_column(connection, "task_memory", "focus_files", "TEXT")
            connection.commit()

    @staticmethod
    def _ensure_column(connection: sqlite3.Connection, table_name: str, column_name: str, column_spec: str) -> None:
        existing_columns = {row[1] for row in connection.execute(f"PRAGMA table_info({table_name})")}
        if column_name in existing_columns:
            return
        connection.execute(f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_spec}")

    async def store_task_result(
        self,
        task_data: TaskState | dict[str, Any],
        *,
        overrides: dict[str, str] | None = None,
    ) -> None:
        if isinstance(task_data, TaskState):
            files_modified = self._extract_files_modified(task_data)
            payload = {
                "task_id": task_data.task_id,
                "goal": task_data.goal,
                "repo": str(task_data.repo_path),
                "solution_summary": task_data.summary or "",
                "success": task_data.status.value == "completed",
                "task_mode": task_data.task_mode,
                "validation_failures": self._extract_validation_failures(task_data),
                "issue_patterns": self._extract_issue_patterns(task_data),
                "focus_files": self._extract_focus_files(task_data),
                "files_modified": files_modified,
            }
        else:
            payload = dict(task_data)
            payload.setdefault("files_modified", [])
            payload.setdefault("solution_summary", "")
            payload.setdefault("success", False)
            payload.setdefault("task_mode", None)
            payload.setdefault("validation_failures", [])
            payload.setdefault("issue_patterns", [])
            payload.setdefault("focus_files", [])

        embedding = await self._embed_or_none(self._task_embedding_text(payload), overrides=overrides)
        if embedding is None:
            return
        with self._lock, self._connect() as connection:
            connection.execute(
                """
                INSERT OR REPLACE INTO task_memory (
                    task_id, goal, repo, solution_summary, success, task_mode, validation_failures,
                    issue_patterns, focus_files, files_modified, timestamp, embedding
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    payload["task_id"],
                    payload["goal"],
                    payload["repo"],
                    payload["solution_summary"],
                    int(bool(payload["success"])),
                    payload["task_mode"],
                    json.dumps(payload["validation_failures"]),
                    json.dumps(payload["issue_patterns"]),
                    json.dumps(payload["focus_files"]),
                    json.dumps(payload["files_modified"]),
                    utc_now().isoformat(),
                    json.dumps(embedding),
                ),
            )
            connection.commit()

    async def retrieve_similar_tasks(
        self,
        query: str,
        *,
        task_mode: str | None = None,
        focus_files: list[str] | None = None,
        validation_failures: list[str] | None = None,
        issue_patterns: list[str] | None = None,
        overrides: dict[str, str] | None = None,
    ) -> list[MemoryHit]:
        query_embedding = await self._embed_or_none(
            self._task_embedding_text(
                {
                    "goal": query,
                    "solution_summary": "",
                    "files_modified": [],
                    "task_mode": task_mode,
                    "validation_failures": validation_failures or [],
                    "issue_patterns": issue_patterns or [],
                    "focus_files": focus_files or [],
                }
            ),
            overrides=overrides,
        )
        if query_embedding is None:
            return []
        with self._lock, self._connect() as connection:
            rows = connection.execute("SELECT * FROM task_memory").fetchall()

        hits: list[MemoryHit] = []
        for row in rows:
            score = self._cosine_similarity(query_embedding, json.loads(row["embedding"]))
            if score <= 0:
                continue
            stored_task_mode = row["task_mode"]
            stored_failures = json.loads(row["validation_failures"]) if row["validation_failures"] else []
            stored_patterns = json.loads(row["issue_patterns"]) if row["issue_patterns"] else []
            stored_focus_files = json.loads(row["focus_files"]) if row["focus_files"] else []
            if task_mode and stored_task_mode == task_mode:
                score += 0.03
            if focus_files:
                overlap = len(set(focus_files) & set(stored_focus_files))
                score += min(0.04, overlap * 0.02)
            if validation_failures:
                overlap = len(set(validation_failures) & set(stored_failures))
                score += min(0.04, overlap * 0.02)
            if issue_patterns:
                overlap = len(set(issue_patterns) & set(stored_patterns))
                score += min(0.04, overlap * 0.02)
            hits.append(
                MemoryHit(
                    task_id=row["task_id"],
                    score=score,
                    summary=row["solution_summary"],
                    source="task_memory",
                    metadata={
                        "task_mode": stored_task_mode,
                        "validation_failures": stored_failures,
                        "issue_patterns": stored_patterns,
                        "focus_files": stored_focus_files,
                        "files_modified": json.loads(row["files_modified"]) if row["files_modified"] else [],
                        "success": bool(row["success"]),
                    },
                )
            )
        hits.sort(key=lambda item: item.score, reverse=True)
        return hits[: self.settings.memory_max_hits]

    async def store_repo_knowledge(
        self,
        repo_data: dict[str, Any],
        *,
        overrides: dict[str, str] | None = None,
    ) -> None:
        repo_name = str(repo_data.get("repo_name", "unknown"))
        embedding = await self._embed_or_none(self._repo_embedding_text(repo_data), overrides=overrides)
        if embedding is None:
            return
        with self._lock, self._connect() as connection:
            connection.execute(
                """
                INSERT OR REPLACE INTO repo_knowledge (repo_name, payload, embedding, updated_at)
                VALUES (?, ?, ?, ?)
                """,
                (
                    repo_name,
                    json.dumps(repo_data),
                    json.dumps(embedding),
                    utc_now().isoformat(),
                ),
            )
            connection.commit()

    async def retrieve_patterns(
        self,
        error_signature: str,
        *,
        overrides: dict[str, str] | None = None,
    ) -> list[str]:
        if not error_signature:
            return []
        query_embedding = await self._embed_or_none(error_signature, overrides=overrides)
        if query_embedding is None:
            return []
        with self._lock, self._connect() as connection:
            rows = connection.execute("SELECT * FROM pattern_memory").fetchall()

        scored_patterns: list[tuple[float, str]] = []
        for row in rows:
            score = self._cosine_similarity(query_embedding, json.loads(row["embedding"]))
            if score > 0:
                scored_patterns.append((score, row["pattern_text"]))
        scored_patterns.sort(key=lambda item: item[0], reverse=True)
        return [pattern for _, pattern in scored_patterns[: self.settings.memory_max_hits]]

    async def record_pattern(
        self,
        error_signature: str,
        pattern_text: str,
        *,
        overrides: dict[str, str] | None = None,
    ) -> None:
        if not error_signature.strip() or not pattern_text.strip():
            return
        embedding = await self._embed_or_none(error_signature, overrides=overrides)
        if embedding is None:
            return
        with self._lock, self._connect() as connection:
            existing = connection.execute(
                "SELECT occurrence_count FROM pattern_memory WHERE error_signature = ?",
                (error_signature,),
            ).fetchone()
            occurrence_count = (existing["occurrence_count"] + 1) if existing else 1
            connection.execute(
                """
                INSERT OR REPLACE INTO pattern_memory (
                    error_signature, pattern_text, occurrence_count, embedding, last_seen
                ) VALUES (?, ?, ?, ?, ?)
                """,
                (
                    error_signature,
                    pattern_text,
                    occurrence_count,
                    json.dumps(embedding),
                    utc_now().isoformat(),
                ),
            )
            connection.commit()

    @staticmethod
    def _extract_validation_failures(task: TaskState) -> list[str]:
        failures: list[str] = []
        for summary in task.validation_history:
            failures.extend(summary.failed_checks)
        return sorted({failure for failure in failures if failure})

    @staticmethod
    def _extract_issue_patterns(task: TaskState) -> list[str]:
        patterns = set(task.reflection_issues)
        if task.focus_issue_inventory is not None:
            patterns.update(issue.issue_type for issue in task.focus_issue_inventory.issues)
        return sorted(pattern for pattern in patterns if pattern)

    @staticmethod
    def _extract_focus_files(task: TaskState) -> list[str]:
        focus_files: set[str] = set(task.retry_focus_files)
        for path in task.repo_context.get("mentioned_files", []):
            normalized = str(path).strip()
            if normalized:
                focus_files.add(normalized)
        if task.focus_issue_inventory is not None:
            focus_files.update(path for path in task.focus_issue_inventory.focus_paths if path)
        return sorted(focus_files)

    @staticmethod
    def _extract_files_modified(task: TaskState) -> list[str]:
        files: set[str] = set()
        for artifact in task.artifacts:
            if artifact.type == "patch" and artifact.path.endswith(".diff"):
                files.add(artifact.path)
        return sorted(files)

    async def _embed_or_none(
        self,
        text: str,
        *,
        overrides: dict[str, str] | None = None,
    ) -> list[float] | None:
        candidate = self._trim_embedding_text(text)
        try:
            return await self.model_router.embed_text(candidate, overrides=overrides)
        except Exception:
            return None

    def _trim_embedding_text(self, text: str) -> str:
        limit = self.settings.memory_embedding_max_chars
        if len(text) <= limit:
            return text
        return text[:limit]

    def _task_embedding_text(self, payload: dict[str, Any]) -> str:
        parts = [
            str(payload.get("goal", "")).strip(),
            str(payload.get("solution_summary", "")).strip(),
            str(payload.get("task_mode", "") or "").strip(),
            " ".join(str(item) for item in payload.get("validation_failures", []) if item),
            " ".join(str(item) for item in payload.get("issue_patterns", []) if item),
            " ".join(str(item) for item in payload.get("focus_files", []) if item),
            " ".join(str(item) for item in payload.get("files_modified", []) if item),
        ]
        return "\n".join(part for part in parts if part)

    def _repo_embedding_text(self, repo_data: dict[str, Any]) -> str:
        payload = {
            "repo_name": repo_data.get("repo_name", "unknown"),
            "mentioned_files": repo_data.get("mentioned_files", []),
            "key_modules": repo_data.get("key_modules", []),
            "important_files": repo_data.get("important_files", []),
            "syntax_findings": repo_data.get("syntax_findings", []),
            "architecture_notes": repo_data.get("architecture_notes", []),
            "entry_points": repo_data.get("entry_points", []),
            "file_previews": repo_data.get("file_previews", []),
        }
        return json.dumps(payload, sort_keys=True)

    @staticmethod
    def _cosine_similarity(vector_a: list[float], vector_b: list[float]) -> float:
        if not vector_a or not vector_b or len(vector_a) != len(vector_b):
            return 0.0
        numerator = sum(a * b for a, b in zip(vector_a, vector_b, strict=False))
        denom_a = math.sqrt(sum(a * a for a in vector_a))
        denom_b = math.sqrt(sum(b * b for b in vector_b))
        if not denom_a or not denom_b:
            return 0.0
        return numerator / (denom_a * denom_b)
