from __future__ import annotations

import json
from collections import defaultdict

from sqlalchemy import text

from app.modules.shared.db import get_engine


class StoryContextSchemaRepository:
    def ensure_tables(self) -> None:
        with get_engine().connect() as conn:
            conn.execute(
                text(
                    """
                    CREATE TABLE IF NOT EXISTS story_records (
                        story_id VARCHAR(128) PRIMARY KEY,
                        project_id VARCHAR(512) NOT NULL,
                        title TEXT NOT NULL,
                        status VARCHAR(64) NOT NULL DEFAULT 'draft',
                        owner VARCHAR(256) NULL,
                        metadata_json JSONB NOT NULL DEFAULT '{}'::jsonb,
                        baseline_commit_sha VARCHAR(128) NULL,
                        snapshot_id VARCHAR(128) NULL,
                        created_by VARCHAR(256) NULL,
                        updated_by VARCHAR(256) NULL,
                        created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
                        updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
                    )
                    """
                )
            )
            conn.execute(
                text(
                    """
                    CREATE TABLE IF NOT EXISTS story_artifacts (
                        id BIGSERIAL PRIMARY KEY,
                        story_id VARCHAR(128) NOT NULL,
                        artifact_type VARCHAR(64) NOT NULL,
                        revision INTEGER NOT NULL,
                        content TEXT NOT NULL,
                        content_hash VARCHAR(128) NULL,
                        rag_session_id VARCHAR(64) NULL,
                        created_by VARCHAR(256) NULL,
                        created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
                        artifact_role VARCHAR(64) NULL,
                        doc_id TEXT NULL,
                        doc_version TEXT NULL,
                        path TEXT NULL,
                        section TEXT NULL,
                        chunk_id TEXT NULL,
                        change_type VARCHAR(32) NULL,
                        summary TEXT NULL,
                        source_ref TEXT NULL,
                        session_id VARCHAR(128) NULL,
                        CONSTRAINT fk_story_artifacts_story
                            FOREIGN KEY (story_id)
                            REFERENCES story_records(story_id)
                            ON DELETE CASCADE,
                        CONSTRAINT uq_story_artifact_revision
                            UNIQUE (story_id, artifact_type, revision)
                    )
                    """
                )
            )
            conn.execute(
                text(
                    """
                    CREATE TABLE IF NOT EXISTS story_links (
                        id BIGSERIAL PRIMARY KEY,
                        story_id VARCHAR(128) NOT NULL,
                        link_type VARCHAR(64) NOT NULL,
                        link_value TEXT NOT NULL,
                        target_ref TEXT NULL,
                        description TEXT NULL,
                        metadata_json JSONB NOT NULL DEFAULT '{}'::jsonb,
                        created_by VARCHAR(256) NULL,
                        created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
                        CONSTRAINT fk_story_links_story
                            FOREIGN KEY (story_id)
                            REFERENCES story_records(story_id)
                            ON DELETE CASCADE
                    )
                    """
                )
            )
            conn.execute(
                text(
                    """
                    CREATE TABLE IF NOT EXISTS session_artifacts (
                        id BIGSERIAL PRIMARY KEY,
                        session_id VARCHAR(128) NOT NULL,
                        project_id VARCHAR(512) NOT NULL,
                        artifact_role VARCHAR(64) NOT NULL,
                        source_ref TEXT NULL,
                        doc_id TEXT NULL,
                        doc_version TEXT NULL,
                        path TEXT NULL,
                        section TEXT NULL,
                        chunk_id TEXT NULL,
                        change_type VARCHAR(32) NULL,
                        summary TEXT NULL,
                        created_by VARCHAR(256) NULL,
                        created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
                        bound_story_id VARCHAR(128) NULL,
                        bound_at TIMESTAMPTZ NULL
                    )
                    """
                )
            )

            # Backward-compatible schema upgrades for existing installations.
            conn.execute(text("ALTER TABLE story_records ADD COLUMN IF NOT EXISTS owner VARCHAR(256) NULL"))
            conn.execute(text("ALTER TABLE story_records ADD COLUMN IF NOT EXISTS metadata_json JSONB NOT NULL DEFAULT '{}'::jsonb"))
            conn.execute(text("ALTER TABLE story_records ADD COLUMN IF NOT EXISTS baseline_commit_sha VARCHAR(128) NULL"))
            conn.execute(text("ALTER TABLE story_records ADD COLUMN IF NOT EXISTS snapshot_id VARCHAR(128) NULL"))
            conn.execute(text("ALTER TABLE story_records ADD COLUMN IF NOT EXISTS created_by VARCHAR(256) NULL"))
            conn.execute(text("ALTER TABLE story_records ADD COLUMN IF NOT EXISTS updated_by VARCHAR(256) NULL"))

            conn.execute(text("ALTER TABLE story_artifacts ADD COLUMN IF NOT EXISTS artifact_role VARCHAR(64) NULL"))
            conn.execute(text("ALTER TABLE story_artifacts ADD COLUMN IF NOT EXISTS doc_id TEXT NULL"))
            conn.execute(text("ALTER TABLE story_artifacts ADD COLUMN IF NOT EXISTS doc_version TEXT NULL"))
            conn.execute(text("ALTER TABLE story_artifacts ADD COLUMN IF NOT EXISTS path TEXT NULL"))
            conn.execute(text("ALTER TABLE story_artifacts ADD COLUMN IF NOT EXISTS section TEXT NULL"))
            conn.execute(text("ALTER TABLE story_artifacts ADD COLUMN IF NOT EXISTS chunk_id TEXT NULL"))
            conn.execute(text("ALTER TABLE story_artifacts ADD COLUMN IF NOT EXISTS change_type VARCHAR(32) NULL"))
            conn.execute(text("ALTER TABLE story_artifacts ADD COLUMN IF NOT EXISTS summary TEXT NULL"))
            conn.execute(text("ALTER TABLE story_artifacts ADD COLUMN IF NOT EXISTS source_ref TEXT NULL"))
            conn.execute(text("ALTER TABLE story_artifacts ADD COLUMN IF NOT EXISTS session_id VARCHAR(128) NULL"))

            conn.execute(text("ALTER TABLE story_links ADD COLUMN IF NOT EXISTS target_ref TEXT NULL"))
            conn.execute(text("ALTER TABLE story_links ADD COLUMN IF NOT EXISTS description TEXT NULL"))
            conn.execute(text("ALTER TABLE story_links ADD COLUMN IF NOT EXISTS metadata_json JSONB NOT NULL DEFAULT '{}'::jsonb"))
            conn.execute(text("ALTER TABLE story_links ADD COLUMN IF NOT EXISTS created_by VARCHAR(256) NULL"))

            conn.execute(text("ALTER TABLE session_artifacts ADD COLUMN IF NOT EXISTS bound_story_id VARCHAR(128) NULL"))
            conn.execute(text("ALTER TABLE session_artifacts ADD COLUMN IF NOT EXISTS bound_at TIMESTAMPTZ NULL"))

            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_story_records_project ON story_records(project_id, updated_at DESC)"))
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_story_artifacts_story_type ON story_artifacts(story_id, artifact_type, revision DESC)"))
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_story_links_story ON story_links(story_id, link_type)"))
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_session_artifacts_session ON session_artifacts(session_id, created_at DESC)"))
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_session_artifacts_project ON session_artifacts(project_id, created_at DESC)"))
            conn.commit()


class StoryContextRepository:
    def record_story_commit(
        self,
        *,
        story_id: str,
        project_id: str,
        title: str,
        commit_sha: str | None,
        branch: str | None,
        changed_files: list[str],
        summary: str,
        actor: str | None,
    ) -> None:
        with get_engine().connect() as conn:
            self._upsert_story_conn(
                conn,
                story_id=story_id,
                project_id=project_id,
                title=title,
                status="in_progress",
                metadata={"source": "repo_webhook"},
                baseline_commit_sha=commit_sha,
                updated_by=actor,
            )
            if commit_sha:
                self._insert_story_link(
                    conn,
                    story_id=story_id,
                    link_type="commit",
                    target_ref=commit_sha,
                    description="Webhook commit",
                    metadata={"project_id": project_id},
                    created_by=actor,
                )
            if branch:
                self._insert_story_link(
                    conn,
                    story_id=story_id,
                    link_type="external",
                    target_ref=branch,
                    description="Webhook branch",
                    metadata={"kind": "branch"},
                    created_by=actor,
                )
            for path in changed_files:
                self._insert_story_link(
                    conn,
                    story_id=story_id,
                    link_type="doc",
                    target_ref=path,
                    description="Changed file in commit",
                    metadata={"commit_sha": commit_sha},
                    created_by=actor,
                )
            revision = self._next_revision(conn, story_id=story_id, artifact_type="commit")
            conn.execute(
                text(
                    """
                    INSERT INTO story_artifacts (
                        story_id,
                        artifact_type,
                        revision,
                        content,
                        content_hash,
                        rag_session_id,
                        created_by,
                        artifact_role,
                        change_type,
                        summary,
                        source_ref
                    ) VALUES (
                        :story_id,
                        :artifact_type,
                        :revision,
                        :content,
                        NULL,
                        NULL,
                        :created_by,
                        :artifact_role,
                        :change_type,
                        :summary,
                        :source_ref
                    )
                    """
                ),
                {
                    "story_id": story_id,
                    "artifact_type": "commit",
                    "revision": revision,
                    "content": summary[:4000] or "commit",
                    "created_by": actor,
                    "artifact_role": "doc_change",
                    "change_type": "linked",
                    "summary": summary[:4000] or "commit",
                    "source_ref": commit_sha or branch or "webhook",
                },
            )
            conn.commit()

    def upsert_story(
        self,
        *,
        story_id: str,
        project_id: str,
        title: str,
        status: str = "draft",
        owner: str | None = None,
        metadata: dict | None = None,
        baseline_commit_sha: str | None = None,
        snapshot_id: str | None = None,
        created_by: str | None = None,
        updated_by: str | None = None,
    ) -> None:
        with get_engine().connect() as conn:
            conn.execute(
                text(
                    """
                    INSERT INTO story_records (
                        story_id,
                        project_id,
                        title,
                        status,
                        owner,
                        metadata_json,
                        baseline_commit_sha,
                        snapshot_id,
                        created_by,
                        updated_by
                    )
                    VALUES (
                        :story_id,
                        :project_id,
                        :title,
                        :status,
                        :owner,
                        CAST(:metadata_json AS JSONB),
                        :baseline_commit_sha,
                        :snapshot_id,
                        :created_by,
                        :updated_by
                    )
                    ON CONFLICT (story_id) DO UPDATE SET
                        project_id = EXCLUDED.project_id,
                        title = EXCLUDED.title,
                        status = EXCLUDED.status,
                        owner = EXCLUDED.owner,
                        metadata_json = EXCLUDED.metadata_json,
                        baseline_commit_sha = COALESCE(EXCLUDED.baseline_commit_sha, story_records.baseline_commit_sha),
                        snapshot_id = COALESCE(EXCLUDED.snapshot_id, story_records.snapshot_id),
                        updated_by = COALESCE(EXCLUDED.updated_by, story_records.updated_by),
                        updated_at = CURRENT_TIMESTAMP
                    """
                ),
                {
                    "story_id": story_id,
                    "project_id": project_id,
                    "title": title,
                    "status": status,
                    "owner": owner,
                    "metadata_json": json.dumps(metadata or {}, ensure_ascii=False),
                    "baseline_commit_sha": baseline_commit_sha,
                    "snapshot_id": snapshot_id,
                    "created_by": created_by,
                    "updated_by": updated_by,
                },
            )
            conn.commit()

    def add_session_artifact(
        self,
        *,
        session_id: str,
        project_id: str,
        artifact_role: str,
        summary: str,
        change_type: str | None = None,
        source_ref: str | None = None,
        doc_id: str | None = None,
        doc_version: str | None = None,
        path: str | None = None,
        section: str | None = None,
        chunk_id: str | None = None,
        created_by: str | None = None,
    ) -> None:
        with get_engine().connect() as conn:
            conn.execute(
                text(
                    """
                    INSERT INTO session_artifacts (
                        session_id,
                        project_id,
                        artifact_role,
                        source_ref,
                        doc_id,
                        doc_version,
                        path,
                        section,
                        chunk_id,
                        change_type,
                        summary,
                        created_by
                    ) VALUES (
                        :session_id,
                        :project_id,
                        :artifact_role,
                        :source_ref,
                        :doc_id,
                        :doc_version,
                        :path,
                        :section,
                        :chunk_id,
                        :change_type,
                        :summary,
                        :created_by
                    )
                    """
                ),
                {
                    "session_id": session_id,
                    "project_id": project_id,
                    "artifact_role": artifact_role,
                    "source_ref": source_ref,
                    "doc_id": doc_id,
                    "doc_version": doc_version,
                    "path": path,
                    "section": section,
                    "chunk_id": chunk_id,
                    "change_type": change_type,
                    "summary": summary,
                    "created_by": created_by,
                },
            )
            conn.commit()

    def bind_session_to_story(
        self,
        *,
        session_id: str,
        story_id: str,
        project_id: str,
        title: str,
        commit_sha: str | None = None,
        branch: str | None = None,
        changed_files: list[str] | None = None,
        actor: str | None = None,
    ) -> dict:
        with get_engine().connect() as conn:
            self._upsert_story_conn(
                conn,
                story_id=story_id,
                project_id=project_id,
                title=title,
                status="in_progress",
                metadata={"bound_session_id": session_id},
                baseline_commit_sha=commit_sha,
                updated_by=actor,
            )

            if commit_sha:
                self._insert_story_link(
                    conn,
                    story_id=story_id,
                    link_type="commit",
                    target_ref=commit_sha,
                    description=f"Bound from session {session_id}",
                    metadata={"session_id": session_id},
                    created_by=actor,
                )
            if branch:
                self._insert_story_link(
                    conn,
                    story_id=story_id,
                    link_type="external",
                    target_ref=branch,
                    description="Source branch",
                    metadata={"kind": "branch"},
                    created_by=actor,
                )
            for path in changed_files or []:
                self._insert_story_link(
                    conn,
                    story_id=story_id,
                    link_type="doc",
                    target_ref=path,
                    description="Changed file from commit",
                    metadata={"session_id": session_id, "commit_sha": commit_sha},
                    created_by=actor,
                )

            rows = conn.execute(
                text(
                    """
                    SELECT id, artifact_role, source_ref, doc_id, doc_version, path, section, chunk_id, change_type, summary, created_by, project_id
                    FROM session_artifacts
                    WHERE session_id = :session_id AND bound_story_id IS NULL
                    ORDER BY id ASC
                    """
                ),
                {"session_id": session_id},
            ).mappings().fetchall()

            revisions = self._load_revisions(conn, story_id)
            migrated = 0
            for row in rows:
                role = str(row["artifact_role"])
                revisions[role] += 1
                summary = str(row["summary"] or "")
                content = summary or "n/a"
                conn.execute(
                    text(
                        """
                        INSERT INTO story_artifacts (
                            story_id,
                            artifact_type,
                            revision,
                            content,
                            content_hash,
                            rag_session_id,
                            created_by,
                            artifact_role,
                            doc_id,
                            doc_version,
                            path,
                            section,
                            chunk_id,
                            change_type,
                            summary,
                            source_ref,
                            session_id
                        ) VALUES (
                            :story_id,
                            :artifact_type,
                            :revision,
                            :content,
                            NULL,
                            :rag_session_id,
                            :created_by,
                            :artifact_role,
                            :doc_id,
                            :doc_version,
                            :path,
                            :section,
                            :chunk_id,
                            :change_type,
                            :summary,
                            :source_ref,
                            :session_id
                        )
                        """
                    ),
                    {
                        "story_id": story_id,
                        "artifact_type": role,
                        "revision": revisions[role],
                        "content": content,
                        "rag_session_id": str(row["project_id"]),
                        "created_by": actor or row["created_by"],
                        "artifact_role": role,
                        "doc_id": row["doc_id"],
                        "doc_version": row["doc_version"],
                        "path": row["path"],
                        "section": row["section"],
                        "chunk_id": row["chunk_id"],
                        "change_type": row["change_type"],
                        "summary": summary,
                        "source_ref": row["source_ref"],
                        "session_id": session_id,
                    },
                )
                conn.execute(
                    text(
                        """
                        UPDATE session_artifacts
                        SET bound_story_id = :story_id, bound_at = CURRENT_TIMESTAMP
                        WHERE id = :id
                        """
                    ),
                    {"story_id": story_id, "id": int(row["id"])},
                )
                migrated += 1

            conn.commit()
            return {"story_id": story_id, "session_id": session_id, "migrated_artifacts": migrated}

    def add_artifact(
        self,
        *,
        story_id: str,
        artifact_type: str,
        revision: int,
        content: str,
        content_hash: str | None = None,
        rag_session_id: str | None = None,
        created_by: str | None = None,
    ) -> None:
        with get_engine().connect() as conn:
            conn.execute(
                text(
                    """
                    INSERT INTO story_artifacts (
                        story_id,
                        artifact_type,
                        revision,
                        content,
                        content_hash,
                        rag_session_id,
                        created_by,
                        artifact_role,
                        summary
                    ) VALUES (
                        :story_id,
                        :artifact_type,
                        :revision,
                        :content,
                        :content_hash,
                        :rag_session_id,
                        :created_by,
                        :artifact_role,
                        :summary
                    )
                    ON CONFLICT (story_id, artifact_type, revision) DO UPDATE SET
                        content = EXCLUDED.content,
                        content_hash = EXCLUDED.content_hash,
                        rag_session_id = EXCLUDED.rag_session_id,
                        created_by = EXCLUDED.created_by,
                        artifact_role = EXCLUDED.artifact_role,
                        summary = EXCLUDED.summary
                    """
                ),
                {
                    "story_id": story_id,
                    "artifact_type": artifact_type,
                    "revision": int(revision),
                    "content": content,
                    "content_hash": content_hash,
                    "rag_session_id": rag_session_id,
                    "created_by": created_by,
                    "artifact_role": artifact_type,
                    "summary": content,
                },
            )
            conn.commit()

    def get_story_context(self, story_id: str) -> dict | None:
        with get_engine().connect() as conn:
            story = conn.execute(
                text(
                    """
                    SELECT story_id, project_id, title, status, owner, metadata_json, baseline_commit_sha, snapshot_id, created_by, updated_by, created_at, updated_at
                    FROM story_records
                    WHERE story_id = :story_id
                    """
                ),
                {"story_id": story_id},
            ).mappings().fetchone()
            if not story:
                return None

            artifacts = conn.execute(
                text(
                    """
                    SELECT artifact_type, artifact_role, revision, content, summary, path, section, chunk_id, change_type, doc_id, doc_version,
                           source_ref, session_id, content_hash, rag_session_id, created_by, created_at
                    FROM story_artifacts
                    WHERE story_id = :story_id
                    ORDER BY artifact_type ASC, revision DESC
                    """
                ),
                {"story_id": story_id},
            ).mappings().fetchall()

            links = conn.execute(
                text(
                    """
                    SELECT link_type,
                           COALESCE(target_ref, link_value) AS target_ref,
                           description,
                           metadata_json,
                           created_by,
                           created_at
                    FROM story_links
                    WHERE story_id = :story_id
                    ORDER BY id ASC
                    """
                ),
                {"story_id": story_id},
            ).mappings().fetchall()

        return {
            "story": dict(story),
            "artifacts": [dict(item) for item in artifacts],
            "links": [dict(item) for item in links],
        }

    def _upsert_story_conn(
        self,
        conn,
        *,
        story_id: str,
        project_id: str,
        title: str,
        status: str,
        metadata: dict,
        baseline_commit_sha: str | None,
        updated_by: str | None,
    ) -> None:
        conn.execute(
            text(
                """
                INSERT INTO story_records (story_id, project_id, title, status, metadata_json, baseline_commit_sha, updated_by)
                VALUES (:story_id, :project_id, :title, :status, CAST(:metadata_json AS JSONB), :baseline_commit_sha, :updated_by)
                ON CONFLICT (story_id) DO UPDATE SET
                    project_id = EXCLUDED.project_id,
                    title = EXCLUDED.title,
                    status = EXCLUDED.status,
                    metadata_json = EXCLUDED.metadata_json,
                    baseline_commit_sha = COALESCE(EXCLUDED.baseline_commit_sha, story_records.baseline_commit_sha),
                    updated_by = COALESCE(EXCLUDED.updated_by, story_records.updated_by),
                    updated_at = CURRENT_TIMESTAMP
                """
            ),
            {
                "story_id": story_id,
                "project_id": project_id,
                "title": title,
                "status": status,
                "metadata_json": json.dumps(metadata, ensure_ascii=False),
                "baseline_commit_sha": baseline_commit_sha,
                "updated_by": updated_by,
            },
        )

    def _insert_story_link(
        self,
        conn,
        *,
        story_id: str,
        link_type: str,
        target_ref: str,
        description: str,
        metadata: dict,
        created_by: str | None,
    ) -> None:
        conn.execute(
            text(
                """
                INSERT INTO story_links (story_id, link_type, link_value, target_ref, description, metadata_json, created_by)
                VALUES (
                    :story_id,
                    :link_type,
                    :link_value,
                    :target_ref,
                    :description,
                    CAST(:metadata_json AS JSONB),
                    :created_by
                )
                """
            ),
            {
                "story_id": story_id,
                "link_type": link_type,
                "link_value": target_ref,
                "target_ref": target_ref,
                "description": description,
                "metadata_json": json.dumps(metadata, ensure_ascii=False),
                "created_by": created_by,
            },
        )

    def _load_revisions(self, conn, story_id: str) -> defaultdict[str, int]:
        rows = conn.execute(
            text(
                """
                SELECT artifact_type, COALESCE(MAX(revision), 0) AS max_revision
                FROM story_artifacts
                WHERE story_id = :story_id
                GROUP BY artifact_type
                """
            ),
            {"story_id": story_id},
        ).mappings().fetchall()
        revision_map: defaultdict[str, int] = defaultdict(int)
        for row in rows:
            revision_map[str(row["artifact_type"])] = int(row["max_revision"])
        return revision_map

    def _next_revision(self, conn, *, story_id: str, artifact_type: str) -> int:
        row = conn.execute(
            text(
                """
                SELECT COALESCE(MAX(revision), 0) AS max_revision
                FROM story_artifacts
                WHERE story_id = :story_id AND artifact_type = :artifact_type
                """
            ),
            {"story_id": story_id, "artifact_type": artifact_type},
        ).mappings().fetchone()
        return int(row["max_revision"]) + 1 if row else 1
