Blame for backend/python-service/storage.py

from __future__ import annotations

import json
import os
import sqlite3
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

from fastapi import HTTPException

from models import (
    ApplyResult,
    ChatConversation,
    ChatConversationDetail,
    ChatConversationSummary,
    ChatMessage,
    ReviewResponse,
    RunDetail,
    RunSummary,
)


PUBLIC_AGENT_NAME = "TabularAgent"


ROOT_DIR = Path(__file__).resolve().parent
DATA_DIR = Path(os.environ.get("MIKE_PLAYBOOK_DATA_DIR", ROOT_DIR / "data"))
DB_PATH = DATA_DIR / "tabular_review.sqlite3"


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def new_id(prefix: str) -> str:
    return f"{prefix}_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:10]}"


def reviewed_docx_name(document_name: str) -> str:
    stem = Path(document_name or "document.docx").stem or "document"
    safe_stem = "".join(
        ch if ch.isalnum() or ch in "._- " else "_" for ch in stem
    ).strip("._ ")
    safe_stem = safe_stem.replace(" ", "_")
    while "__" in safe_stem:
        safe_stem = safe_stem.replace("__", "_")
    return f"{safe_stem or 'document'}_tabular_reviewed.docx"


def short_timestamp(value: str) -> str:
    try:
        dt = datetime.fromisoformat(value)
    except ValueError:
        return value[:16]
    return dt.astimezone(timezone.utc).strftime("%Y-%m-%d %H:%M")


def connection() -> sqlite3.Connection:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def init_db() -> None:
    with connection() as conn:
        conn.executescript(
            """
            create table if not exists workflow_runs (
                run_id text primary key,
                owner_username text not null default 'admin',
                created_at text not null,
                document_name text not null,
                original_docx_path text not null,
                commented_docx_path text,
                document_text text not null,
                playbook_id text not null,
                playbook_name text not null,
                playbook_text text not null,
                max_findings integer not null,
                status text not null,
                result_json text not null,
                apply_report_json text not null,
                model text,
                usage_json text
            );

            create table if not exists chat_conversations (
                conversation_id text primary key,
                owner_username text not null default 'admin',
                title text not null,
                created_at text not null,
                updated_at text not null
            );

            create table if not exists chat_messages (
                message_id text primary key,
                conversation_id text not null,
                owner_username text not null default 'admin',
                created_at text not null,
                role text not null,
                content text not null,
                document_name text,
                foreign key(conversation_id) references chat_conversations(conversation_id)
            );
            """
        )
        _ensure_column(
            conn, "workflow_runs", "owner_username", "text not null default 'admin'"
        )
        _ensure_column(
            conn,
            "chat_conversations",
            "owner_username",
            "text not null default 'admin'",
        )
        _ensure_column(
            conn, "chat_messages", "owner_username", "text not null default 'admin'"
        )


def _ensure_column(
    conn: sqlite3.Connection, table: str, column: str, definition: str
) -> None:
    existing = {
        row["name"] for row in conn.execute(f"pragma table_info({table})").fetchall()
    }
    if column not in existing:
        conn.execute(f"alter table {table} add column {column} {definition}")


def save_workflow_run(
    *,
    owner_username: str,
    run_id: str,
    document_name: str,
    original_docx_path: Path,
    commented_docx_path: Optional[Path],
    document_text: str,
    playbook_id: str,
    playbook_name: str,
    playbook_text: str,
    max_findings: int,
    status: str,
    result: ReviewResponse,
    apply_report: list[ApplyResult],
    model: Optional[str],
    usage: Optional[dict[str, Any]],
) -> None:
    with connection() as conn:
        conn.execute(
            """
            insert or replace into workflow_runs (
                run_id, owner_username, created_at, document_name, original_docx_path, commented_docx_path,
                document_text, playbook_id, playbook_name, playbook_text, max_findings,
                status, result_json, apply_report_json, model, usage_json
            )
            values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                run_id,
                owner_username,
                utc_now(),
                document_name,
                str(original_docx_path),
                str(commented_docx_path) if commented_docx_path else None,
                document_text,
                playbook_id,
                playbook_name,
                playbook_text,
                max_findings,
                status,
                result.model_dump_json(),
                json.dumps(
                    [item.model_dump() for item in apply_report], ensure_ascii=False
                ),
                model,
                json.dumps(usage, ensure_ascii=False) if usage else None,
            ),
        )


def _summary_from_row(row: sqlite3.Row) -> RunSummary:
    result = json.loads(row["result_json"])
    apply_report = json.loads(row["apply_report_json"])
    findings = result.get("findings") if isinstance(result, dict) else []
    applied_count = sum(1 for item in apply_report if item.get("status") == "applied")
    return RunSummary(
        run_id=row["run_id"],
        created_at=row["created_at"],
        created_short=short_timestamp(row["created_at"]),
        document_name=row["document_name"],
        commented_docx_name=reviewed_docx_name(row["document_name"]),
        playbook_id=row["playbook_id"],
        playbook_name=row["playbook_name"],
        status=row["status"],
        finding_count=len(findings) if isinstance(findings, list) else 0,
        applied_count=applied_count,
        model=PUBLIC_AGENT_NAME,
    )


def list_runs(
    owner_username: str, limit: int = 50, playbook_id: str | None = None
) -> list[RunSummary]:
    with connection() as conn:
        if playbook_id:
            rows = conn.execute(
                "select * from workflow_runs where owner_username = ? and playbook_id = ? order by created_at desc limit ?",
                (owner_username, playbook_id, limit),
            ).fetchall()
        else:
            rows = conn.execute(
                "select * from workflow_runs where owner_username = ? order by created_at desc limit ?",
                (owner_username, limit),
            ).fetchall()
    return [_summary_from_row(row) for row in rows]


def get_run(run_id: str, owner_username: str) -> RunDetail:
    with connection() as conn:
        row = conn.execute(
            "select * from workflow_runs where run_id = ? and owner_username = ?",
            (run_id, owner_username),
        ).fetchone()
    if not row:
        raise HTTPException(status_code=404, detail=f"Workflow run not found: {run_id}")
    summary = _summary_from_row(row)
    usage = json.loads(row["usage_json"]) if row["usage_json"] else None
    commented_path = (
        Path(row["commented_docx_path"]) if row["commented_docx_path"] else None
    )
    return RunDetail(
        **summary.model_dump(),
        document_text=row["document_text"],
        playbook_text=row["playbook_text"],
        result=ReviewResponse.model_validate_json(row["result_json"]),
        apply_report=[
            ApplyResult.model_validate(item)
            for item in json.loads(row["apply_report_json"])
        ],
        usage=usage,
        commented_docx_available=bool(commented_path and commented_path.exists()),
    )


def get_commented_docx_path(run_id: str, owner_username: str) -> Path:
    with connection() as conn:
        row = conn.execute(
            "select commented_docx_path from workflow_runs where run_id = ? and owner_username = ?",
            (run_id, owner_username),
        ).fetchone()
    if not row or not row["commented_docx_path"]:
        raise HTTPException(
            status_code=404, detail="Commented DOCX is not available for this run."
        )
    path = Path(row["commented_docx_path"])
    if not path.exists():
        raise HTTPException(status_code=404, detail="Commented DOCX file is missing.")
    return path


def save_chat_message(
    *,
    owner_username: str,
    conversation_id: Optional[str],
    role: str,
    content: str,
    document_name: Optional[str] = None,
) -> str:
    now = utc_now()
    conv_id = conversation_id or new_id("conv")
    title = content.strip().splitlines()[0][:80] or "New conversation"
    message_id = new_id("msg")
    with connection() as conn:
        existing = conn.execute(
            "select conversation_id from chat_conversations where conversation_id = ? and owner_username = ?",
            (conv_id, owner_username),
        ).fetchone()
        if not existing:
            conn.execute(
                "insert into chat_conversations (conversation_id, owner_username, title, created_at, updated_at) values (?, ?, ?, ?, ?)",
                (conv_id, owner_username, title, now, now),
            )
        else:
            conn.execute(
                "update chat_conversations set updated_at = ? where conversation_id = ? and owner_username = ?",
                (now, conv_id, owner_username),
            )
        conn.execute(
            """
            insert into chat_messages (
                message_id, conversation_id, owner_username, created_at, role, content, document_name
            )
            values (?, ?, ?, ?, ?, ?, ?)
            """,
            (
                message_id,
                conv_id,
                owner_username,
                now,
                role,
                content,
                document_name,
            ),
        )
    return conv_id


def list_conversations(
    owner_username: str, limit: int = 20
) -> list[ChatConversationSummary]:
    with connection() as conn:
        rows = conn.execute(
            """
            select conversation_id, title, created_at, updated_at
            from chat_conversations
            where owner_username = ?
            order by updated_at desc
            limit ?
            """,
            (owner_username, limit),
        ).fetchall()
    return [
        ChatConversationSummary(
            conversation_id=row["conversation_id"],
            title=row["title"],
            created_at=row["created_at"],
            updated_at=row["updated_at"],
        )
        for row in rows
    ]


def get_conversation(
    conversation_id: str, owner_username: str
) -> ChatConversationDetail:
    with connection() as conn:
        row = conn.execute(
            """
            select conversation_id, owner_username, title, created_at, updated_at
            from chat_conversations
            where conversation_id = ? and owner_username = ?
            """,
            (conversation_id, owner_username),
        ).fetchone()
        if not row:
            raise HTTPException(
                status_code=404, detail=f"Conversation not found: {conversation_id}"
            )
        messages = conn.execute(
            """
            select message_id, created_at, role, content, document_name
            from chat_messages
            where conversation_id = ? and owner_username = ?
            order by created_at asc
            """,
            (conversation_id, owner_username),
        ).fetchall()
    return ChatConversationDetail(
        conversation_id=row["conversation_id"],
        owner_username=row["owner_username"],
        title=row["title"],
        created_at=row["created_at"],
        updated_at=row["updated_at"],
        messages=[
            ChatMessage(
                id=message["message_id"],
                role=message["role"],
                content=message["content"],
                created_at=message["created_at"],
                document_name=message["document_name"],
            )
            for message in messages
        ],
    )


def run_to_markdown(run_id: str, owner_username: str) -> str:
    run = get_run(run_id, owner_username)
    lines = [
        f"# Review Run: {run.document_name}",
        "",
        "## Metadata",
        "",
        f"- Run ID: `{run.run_id}`",
        f"- Created: `{run.created_short}`",
        f"- Playbook: `{run.playbook_name}` (`{run.playbook_id}`)",
        f"- Status: `{run.status}`",
        f"- Agent: `{PUBLIC_AGENT_NAME}`",
        f"- Findings: `{run.finding_count}`",
        f"- Comments applied: `{run.applied_count}`",
        "",
        "## Findings",
        "",
    ]
    for index, finding in enumerate(run.result.findings, start=1):
        apply_result = next(
            (item for item in run.apply_report if item.id == finding.id), None
        )
        lines.extend(
            [
                f"### {index}. {finding.clause_type} - {finding.risk_level}",
                "",
                f"**Issue:** {finding.issue}",
                "",
                f"**Comment inserted:** {finding.comment}",
                "",
                f"**Internal reasoning:**\n- **Clause/Section:** {finding.review_reasoning.get('clause_section', 'N/A') if isinstance(finding.review_reasoning, dict) else finding.review_reasoning}\n- **Violated Standard:** {finding.review_reasoning.get('violated_standard', 'N/A') if isinstance(finding.review_reasoning, dict) else 'N/A'}\n- **Zamil Exposure:** {finding.review_reasoning.get('zamil_exposure', 'N/A') if isinstance(finding.review_reasoning, dict) else 'N/A'}",
                "",
                f"**Confidence:** {finding.confidence}",
                "",
                f"**Apply status:** {apply_result.status if apply_result else 'not_applied'}",
                "",
                "**Target text:**",
                "",
                "```text",
                finding.target_text,
                "```",
                "",
            ]
        )
    if not run.result.findings:
        lines.append("No findings were returned.")
        lines.append("")
    lines.extend(
        [
            "## Apply Report",
            "",
            "```json",
            json.dumps(
                [item.model_dump() for item in run.apply_report],
                indent=2,
                ensure_ascii=False,
            ),
            "```",
            "",
            "## Document Text",
            "",
            "```text",
            run.document_text,
            "```",
            "",
            "## Playbook Text",
            "",
            "```yaml",
            run.playbook_text,
            "```",
            "",
            "## Raw Result JSON",
            "",
            "```json",
            run.result.model_dump_json(indent=2),
            "```",
            "",
        ]
    )
    return "\n".join(lines)