#!/usr/bin/env python3
import argparse
import datetime as dt
import json
import os
import re
import sqlite3
from typing import Any, Dict, Iterable, List, Optional

import mm_users


def open_db(db_path: str, *, readonly: bool = False, timeout: float = 5.0) -> sqlite3.Connection:
    """Open sqlite index with busy timeout; readonly avoids write locks on shared DB."""
    if readonly:
        abs_path = os.path.abspath(db_path)
        conn = sqlite3.connect(f"file:{abs_path}?mode=ro", uri=True, timeout=timeout)
    else:
        conn = sqlite3.connect(db_path, timeout=timeout)
    conn.row_factory = sqlite3.Row
    conn.execute(f"PRAGMA busy_timeout={int(timeout * 1000)}")
    return conn


def iter_result_files(root: str) -> Iterable[str]:
    for dirpath, _dirnames, filenames in os.walk(root):
        for name in filenames:
            if not name.endswith(".json"):
                continue
            yield os.path.join(dirpath, name)


def parse_iso_date(s: str) -> dt.datetime:
    # Accept YYYY-MM-DD or full ISO with time; normalize to UTC
    try:
        if "T" in s:
            return dt.datetime.fromisoformat(s.replace("Z", "+00:00")).astimezone(dt.timezone.utc)
        return dt.datetime.strptime(s, "%Y-%m-%d").replace(tzinfo=dt.timezone.utc)
    except ValueError as e:
        raise argparse.ArgumentTypeError(f"invalid date '{s}': {e}") from e


def parse_post_created(s: str) -> dt.datetime:
    # Exporter uses '%Y-%m-%dT%H:%M:%SZ'
    return dt.datetime.strptime(s, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=dt.timezone.utc)


def make_matcher(pattern: str, use_regex: bool, ignore_case: bool):
    if use_regex:
        flags = re.IGNORECASE if ignore_case else 0
        rx = re.compile(pattern, flags)

        def match(text: str) -> bool:
            return bool(rx.search(text))

        return match

    needle = pattern.lower() if ignore_case else pattern

    def match(text: str) -> bool:
        hay = text.lower() if ignore_case else text
        return needle in hay

    return match


def build_fts_query(query: str) -> Optional[str]:
    # Tokenize user query and run prefix search by tokens.
    tokens = re.findall(r"\w+", query, flags=re.UNICODE)
    if not tokens:
        return None
    return " AND ".join(f"{tok}*" for tok in tokens)


def _escape_like(s: str) -> str:
    return s.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")


def _post_row_to_match(row: sqlite3.Row) -> Dict[str, Any]:
    files: Optional[List[str]] = None
    if row["files_json"]:
        try:
            maybe_files = json.loads(row["files_json"])
            if isinstance(maybe_files, list):
                files = maybe_files
        except json.JSONDecodeError:
            files = None
    post: Dict[str, Any] = {
        "id": row["post_id"],
        "idx": row["idx"],
        "created": row["created"].replace("+00:00", "Z"),
        "username": row["username"],
        "message": row["message"],
    }
    if files:
        post["files"] = files
    return {
        "file": row["source_path"],
        "channel": row["channel_display_name"] or row["channel_name"],
        "team": row["team"],
        "post": post,
    }


def search_tag(
    db_path: str,
    keyword: str,
    after: dt.datetime,
    before: dt.datetime,
    *,
    user: Optional[str] = None,
    channel: Optional[str] = None,
    team: Optional[str] = None,
    limit: Optional[int] = 50,
) -> Iterable[Dict[str, Any]]:
    """Substring search for #keyword within a required created date range (newest first).

    Uses plain SQL LIKE — no FTS/rank. Intended for tags like plan/report over a day/week.
    """
    raw = (keyword or "").strip()
    if raw.startswith("#"):
        raw = raw[1:].strip()
    if not raw:
        raise ValueError("keyword must be non-empty (e.g. 'plan' or 'report')")
    if after is None or before is None:
        raise ValueError("after and before are required")
    if after > before:
        raise ValueError("after must be <= before")

    tag = f"#{raw}"
    needle = f"%{_escape_like(tag)}%"

    conn = open_db(db_path, readonly=True)
    try:
        filters = [
            "p.created >= ?",
            "p.created <= ?",
            "p.message LIKE ? ESCAPE '\\'",
        ]
        params: List[Any] = [after.isoformat(), before.isoformat(), needle]

        if user:
            filters.append("p.username = ?")
            params.append(user)
        if channel:
            filters.append("(p.channel_display_name = ? OR p.channel_name = ?)")
            params.extend([channel, channel])
        if team:
            filters.append("p.team = ?")
            params.append(team)

        sql = f"""
            SELECT
                p.post_id, p.idx, p.created, p.username, p.message, p.files_json,
                p.channel_display_name, p.channel_name, p.team, p.source_path
            FROM posts p
            WHERE {' AND '.join(filters)}
            ORDER BY p.created DESC
        """
        if limit is not None:
            sql += " LIMIT ?"
            params.append(limit)

        for row in conn.execute(sql, params).fetchall():
            yield _post_row_to_match(row)
    finally:
        conn.close()


def search_index(
    db_path: str,
    query: str,
    user: Optional[str],
    channel: Optional[str],
    team: Optional[str],
    after: Optional[dt.datetime],
    before: Optional[dt.datetime],
    limit: Optional[int],
    ignore_case: bool,
) -> Iterable[Dict[str, Any]]:
    conn = open_db(db_path, readonly=True)
    try:
        filters = []
        params: List[Any] = []

        if user:
            filters.append("p.username = ?")
            params.append(user)
        if channel:
            filters.append("(p.channel_display_name = ? OR p.channel_name = ?)")
            params.extend([channel, channel])
        if team:
            filters.append("p.team = ?")
            params.append(team)
        if after:
            filters.append("p.created >= ?")
            params.append(after.isoformat())
        if before:
            filters.append("p.created <= ?")
            params.append(before.isoformat())

        where_filters = f"AND {' AND '.join(filters)}" if filters else ""
        fts_q = build_fts_query(query)
        if fts_q:
            # Primary: FTS relevance (rank). Tie-break: newest first.
            sql = f"""
                SELECT
                    p.post_id, p.idx, p.created, p.username, p.message, p.files_json,
                    p.channel_display_name, p.channel_name, p.team, p.source_path
                FROM posts_fts fts
                JOIN posts p ON p.post_id = fts.post_id
                WHERE fts.message MATCH ?
                {where_filters}
                ORDER BY rank, p.created DESC
            """
            params = [fts_q] + params
        else:
            cmp_col = "LOWER(p.message)" if ignore_case else "p.message"
            needle = query.lower() if ignore_case else query
            sql = f"""
                SELECT
                    p.post_id, p.idx, p.created, p.username, p.message, p.files_json,
                    p.channel_display_name, p.channel_name, p.team, p.source_path
                FROM posts p
                WHERE {cmp_col} LIKE ?
                {where_filters}
                ORDER BY p.created DESC
            """
            params = [f"%{needle}%"] + params

        if limit is not None:
            sql += " LIMIT ?"
            params.append(limit)

        rows = conn.execute(sql, params).fetchall()
        for row in rows:
            yield _post_row_to_match(row)
    finally:
        conn.close()


def search_attachments_index(
    db_path: str,
    query: str,
    user: Optional[str],
    channel: Optional[str],
    team: Optional[str],
    after: Optional[dt.datetime],
    before: Optional[dt.datetime],
    limit: Optional[int],
    mode: str,
) -> Iterable[Dict[str, Any]]:
    """
    Search in attachments index.

    mode: "name" to search by file name, "content" to search by text content.
    """
    conn = open_db(db_path, readonly=True)
    try:
        filters = []
        params: List[Any] = []

        if user:
            filters.append("p.username = ?")
            params.append(user)
        if channel:
            filters.append("(p.channel_display_name = ? OR p.channel_name = ?)")
            params.extend([channel, channel])
        if team:
            filters.append("p.team = ?")
            params.append(team)
        if after:
            filters.append("p.created >= ?")
            params.append(after.isoformat())
        if before:
            filters.append("p.created <= ?")
            params.append(before.isoformat())

        where_filters = f"AND {' AND '.join(filters)}" if filters else ""

        fts_column = "filename" if mode == "name" else "content"
        fts_q = build_fts_query(query)
        if not fts_q:
            return []

        sql = f"""
            SELECT
                p.post_id,
                p.idx,
                p.created,
                p.username,
                p.message,
                p.files_json,
                p.channel_display_name,
                p.channel_name,
                p.team,
                p.source_path,
                a.filename AS attachment_filename,
                a.filepath AS attachment_path,
                a.is_text AS attachment_is_text
            FROM attachments_fts aft
            JOIN attachments a ON a.id = aft.attachment_id
            JOIN posts p ON p.post_id = a.post_id
            WHERE aft.{fts_column} MATCH ?
            {where_filters}
            ORDER BY rank, p.created DESC
        """

        params = [fts_q] + params
        if limit is not None:
            sql += " LIMIT ?"
            params.append(limit)

        rows = conn.execute(sql, params).fetchall()
        for row in rows:
            files: Optional[List[str]] = None
            if row["files_json"]:
                try:
                    maybe_files = json.loads(row["files_json"])
                    if isinstance(maybe_files, list):
                        files = maybe_files
                except json.JSONDecodeError:
                    files = None
            post: Dict[str, Any] = {
                "id": row["post_id"],
                "idx": row["idx"],
                "created": row["created"].replace("+00:00", "Z"),
                "username": row["username"],
                "message": row["message"],
            }
            if files:
                post["files"] = files

            attachment = {
                "filename": row["attachment_filename"],
                "path": row["attachment_path"],
                "is_text": bool(row["attachment_is_text"]),
            }

            yield {
                "file": row["source_path"],
                "channel": row["channel_display_name"] or row["channel_name"],
                "team": row["team"],
                "post": post,
                "attachment": attachment,
            }
    finally:
        conn.close()


def search_file(
    path: str,
    matcher,
    user: Optional[str],
    channel: Optional[str],
    team: Optional[str],
    after: Optional[dt.datetime],
    before: Optional[dt.datetime],
) -> Iterable[Dict[str, Any]]:
    try:
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
    except (OSError, json.JSONDecodeError):
        return []

    ch = data.get("channel", {}) or {}
    posts: List[Dict[str, Any]] = data.get("posts") or []

    channel_name = ch.get("display_name") or ch.get("name") or ""
    team_name = ch.get("team") or ""

    if channel and channel not in (channel_name, ch.get("name", "")):
        return []
    if team and team != team_name:
        return []

    for p in posts:
        msg = p.get("message") or ""
        username = p.get("username") or ""
        created_raw = p.get("created")
        if not isinstance(created_raw, str):
            continue
        try:
            created = parse_post_created(created_raw)
        except ValueError:
            continue

        if user and username != user:
            continue
        if after and created < after:
            continue
        if before and created > before:
            continue
        if not matcher(msg):
            continue

        yield {
            "file": path,
            "channel": channel_name,
            "team": team_name,
            "post": p,
        }


def print_human(matches: Iterable[Dict[str, Any]], limit: Optional[int], show_file: bool) -> None:
    count = 0
    for m in matches:
        if limit is not None and count >= limit:
            break
        p = m["post"]
        created = p.get("created", "")
        username = p.get("username", "")
        msg = (p.get("message") or "").replace("\n", "\\n")
        attachment = m.get("attachment")
        attachment_suffix = ""
        if isinstance(attachment, dict) and attachment.get("filename"):
            attachment_suffix = f" [file: {attachment['filename']}]"
        if show_file:
            file_rel = os.path.relpath(m["file"])
            idx = p.get("idx", "?")
            prefix = f"{file_rel}:{idx}: "
        else:
            prefix = ""
        channel = m.get("channel") or "?"
        team = m.get("team") or ""
        team_prefix = f"[{team}]" if team else ""
        print(f"{prefix}{team_prefix}[{channel}] {created} {username}: {msg}{attachment_suffix}")
        count += 1


def print_json(matches: Iterable[Dict[str, Any]], limit: Optional[int]) -> None:
    out: List[Dict[str, Any]] = []
    count = 0
    for m in matches:
        if limit is not None and count >= limit:
            break
        out.append(m)
        count += 1
    json.dump(out, fp=os.sys.stdout, ensure_ascii=False, indent=2)
    print()


def main(argv: Optional[List[str]] = None) -> int:
    ap = argparse.ArgumentParser(
        description="Search through Mattermost export JSON files produced by mattermost-dl.py.",
    )
    ap.add_argument("query", nargs="?", help="Search string or regex (see --regex).")
    ap.add_argument(
        "-r",
        "--root",
        default="results",
        help="Root directory with export results (default: results).",
    )
    ap.add_argument("--user", help="Filter by username (exact match).")
    ap.add_argument("--channel", help="Filter by channel display_name or name (exact match).")
    ap.add_argument("--team", help="Filter by team name (exact match).")
    ap.add_argument("--after", type=parse_iso_date, help="Only messages on/after this date (YYYY-MM-DD or ISO).")
    ap.add_argument("--before", type=parse_iso_date, help="Only messages on/before this date (YYYY-MM-DD or ISO).")
    ap.add_argument("--regex", action="store_true", help="Treat query as regular expression.")
    group = ap.add_mutually_exclusive_group()
    group.add_argument(
        "--file-name",
        action="store_true",
        help="Search by attachment file name instead of message text (requires sqlite index).",
    )
    group.add_argument(
        "--file-content",
        action="store_true",
        help="Search by attachment text content instead of message text (requires sqlite index; only UTF-8 text files are indexed).",
    )
    ap.add_argument(
        "-i",
        "--ignore-case",
        action="store_true",
        help="Case-insensitive search (default for substring search).",
    )
    ap.add_argument("--json", action="store_true", help="Output matches as JSON array (for jq, etc.).")
    ap.add_argument(
        "-n",
        "--limit",
        type=int,
        help="Maximum number of matches to output.",
    )
    ap.add_argument("--db", default="mattermost_index.sqlite", help="SQLite index database path.")
    ap.add_argument(
        "--no-index",
        action="store_true",
        help="Do not use sqlite index even if DB exists.",
    )
    ap.add_argument(
        "--show-file",
        action="store_true",
        help="Include source filename and per-file index in human output.",
    )
    user_group = ap.add_mutually_exclusive_group()
    user_group.add_argument(
        "--resolve-user",
        metavar="REF",
        help="Resolve display name / FIO / nickname to username (uses sqlite user directory).",
    )
    user_group.add_argument(
        "--search-users",
        metavar="QUERY",
        help="Search users by name or username substring.",
    )
    ap.add_argument(
        "--refresh-users",
        action="store_true",
        help="Refresh user directory from users.json, optional API, and post usernames.",
    )
    ap.add_argument(
        "--from-api",
        action="store_true",
        help="With --refresh-users, also fetch profiles from Mattermost API (config.json).",
    )
    ap.add_argument(
        "--config",
        default=os.environ.get("MM_CONFIG_PATH", "config.json"),
        help="Mattermost config path for --from-api (default: MM_CONFIG_PATH or config.json).",
    )

    args = ap.parse_args(argv)

    if args.refresh_users:
        result = mm_users.refresh_user_directory(
            args.db,
            args.root,
            config_path=args.config,
            from_api=args.from_api,
        )
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return 0

    if args.resolve_user:
        result = mm_users.resolve_user(args.db, args.resolve_user)
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return 0

    if args.search_users:
        matches = mm_users.search_users(args.db, args.search_users, limit=args.limit or 10)
        print(json.dumps(matches, ensure_ascii=False, indent=2))
        return 0

    if not args.query:
        ap.error("query is required unless using --resolve-user, --search-users, or --refresh-users")

    # Default ignore-case for non-regex search unless explicitly disabled
    ignore_case = args.ignore_case or not args.regex
    matcher = make_matcher(args.query, args.regex, ignore_case)

    if (args.file_name or args.file_content) and args.regex:
        raise SystemExit("--file-name/--file-content do not support --regex yet.")

    use_index = (not args.no_index) and os.path.exists(args.db)
    if args.file_name or args.file_content:
        if not use_index:
            raise SystemExit(
                "Attachment search (--file-name/--file-content) requires sqlite index. "
                "Create it with mm_index.py and do not use --no-index."
            )
        mode = "name" if args.file_name else "content"
        matches_iter = search_attachments_index(
            db_path=args.db,
            query=args.query,
            user=args.user,
            channel=args.channel,
            team=args.team,
            after=args.after,
            before=args.before,
            limit=args.limit,
            mode=mode,
        )
        output_limit = None
    elif use_index and not args.regex:
        matches_iter = search_index(
            db_path=args.db,
            query=args.query,
            user=args.user,
            channel=args.channel,
            team=args.team,
            after=args.after,
            before=args.before,
            limit=args.limit,
            ignore_case=ignore_case,
        )
        # Limit already applied in SQL.
        output_limit = None
    else:
        files = list(iter_result_files(args.root))
        matches_iter = (
            m
            for path in files
            for m in search_file(
                path=path,
                matcher=matcher,
                user=args.user,
                channel=args.channel,
                team=args.team,
                after=args.after,
                before=args.before,
            )
        )
        output_limit = args.limit

    if args.json:
        print_json(matches_iter, output_limit)
    else:
        print_human(matches_iter, output_limit, show_file=args.show_file)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())