"""
apitest-mcp — MCP server for API migration testing.

Tools:
  - scan_repo_config   : scan repo to find PG connection string
  - read_api_code      : scan repo to understand an API endpoint's params/tables
  - query_pg           : run a SELECT on PostgreSQL
  - call_api           : call an API endpoint
  - validate_db        : run a SELECT to validate DB state after API call
"""

import os
import re
import json
import glob
import subprocess
import urllib.request
import urllib.parse
import urllib.error
from pathlib import Path

import psycopg2
import psycopg2.extras
from mcp.server.fastmcp import FastMCP
from dotenv import load_dotenv

load_dotenv()

EMP_ID = os.getenv("EMP_ID", "")
AK = os.getenv("AK", "")
REPO_PATH = os.getenv("REPO_PATH", "")

mcp = FastMCP("apitest-mcp")

# ── internal helpers ──────────────────────────────────────────────────────────

_pg_conn_cache: dict = {}


def _find_pg_connection(repo_path: str) -> dict:
    """
    Scan a repo for PostgreSQL connection strings.
    Supports PHP (Yii config), Go (.env / yaml / go files), generic .env files.
    Returns {"host": ..., "port": ..., "dbname": ..., "user": ..., "password": ...}
    """
    if repo_path in _pg_conn_cache:
        return _pg_conn_cache[repo_path]

    found = {}

    # ── 1. PHP Yii: config/main.php ──
    for cfg in Path(repo_path).rglob("main.php"):
        text = cfg.read_text(errors="ignore")
        # look for pgsql:host=...;port=...;dbname=...
        m = re.search(r"pgsql:host=([^;'\"]+);port=(\d+);dbname=([^'\"]+)", text)
        if m:
            found["host"] = m.group(1).strip()
            found["port"] = int(m.group(2))
            found["dbname"] = m.group(3).strip()
        um = re.search(r"'username'\s*=>\s*'([^']+)'", text)
        if um:
            found["user"] = um.group(1)
        pm = re.search(r"'password'\s*=>\s*'([^']+)'", text)
        if pm:
            found["password"] = pm.group(1)
        if found.get("host"):
            break

    # ── 2. .env in repo root ──
    if not found.get("host"):
        env_file = Path(repo_path) / ".env"
        if env_file.exists():
            text = env_file.read_text(errors="ignore")
            m = re.search(r"(?:DATABASE_URL|PG_URL|POSTGRES_URL)\s*=\s*postgresql://([^:]+):([^@]+)@([^:/]+):?(\d*)/(\S+)", text)
            if m:
                found["user"] = m.group(1)
                found["password"] = m.group(2)
                found["host"] = m.group(3)
                found["port"] = int(m.group(4)) if m.group(4) else 5432
                found["dbname"] = m.group(5)

    # ── 3. Go secrets / yaml ──
    if not found.get("host"):
        for f in list(Path(repo_path).rglob("*.yaml")) + list(Path(repo_path).rglob("*.yml")):
            text = f.read_text(errors="ignore")
            m = re.search(r"host:\s*(\S+).*?port:\s*(\d+).*?dbname:\s*(\S+).*?user:\s*(\S+).*?password:\s*(\S+)", text, re.DOTALL)
            if m:
                found = {"host": m.group(1), "port": int(m.group(2)),
                         "dbname": m.group(3), "user": m.group(4), "password": m.group(5)}
                break

    # ── 4. Go source: pgx / pq connection strings ──
    if not found.get("host"):
        for f in Path(repo_path).rglob("*.go"):
            text = f.read_text(errors="ignore")
            m = re.search(r"host=(\S+)\s+port=(\d+)\s+dbname=(\S+)\s+user=(\S+)\s+password=(\S+)", text)
            if m:
                found = {"host": m.group(1), "port": int(m.group(2)),
                         "dbname": m.group(3), "user": m.group(4), "password": m.group(5)}
                break

    _pg_conn_cache[repo_path] = found
    return found


def _get_pg_conn(repo_path: str):
    creds = _find_pg_connection(repo_path)
    if not creds.get("host"):
        raise ValueError(f"Could not find PostgreSQL credentials in repo: {repo_path}")
    return psycopg2.connect(
        host=creds["host"],
        port=creds.get("port", 5432),
        dbname=creds["dbname"],
        user=creds["user"],
        password=creds["password"],
        connect_timeout=10,
    )


# ── MCP tools ─────────────────────────────────────────────────────────────────


@mcp.tool()
def scan_repo_config(repo_path: str = "") -> str:
    """
    Scan a repository to find PostgreSQL connection details.
    Returns the discovered host, port, dbname (passwords redacted).
    Use this first to confirm the DB can be reached before running queries.

    Args:
        repo_path: Absolute path to the repository root. Defaults to REPO_PATH env var.
    """
    path = repo_path or REPO_PATH
    if not path:
        return "ERROR: repo_path not provided and REPO_PATH env var is not set."
    creds = _find_pg_connection(path)
    if not creds:
        return f"No PostgreSQL credentials found in {path}"
    safe = {k: v for k, v in creds.items() if k != "password"}
    safe["password"] = "***"
    return json.dumps(safe, indent=2)


@mcp.tool()
def read_api_code(api_url: str, repo_path: str = "") -> str:
    """
    Given an API URL, scan the repository to find the handler code for that endpoint.
    Returns relevant source code snippets: route definition, controller action, model queries.
    Use this to understand what params the API expects and which DB tables it touches.

    Args:
        api_url: Full URL of the API endpoint, e.g. http://localhost:8080/erp/preovp/preovpdata
        repo_path: Absolute path to the repository root. Defaults to REPO_PATH env var.
    """
    path = repo_path or REPO_PATH
    if not path:
        return "ERROR: repo_path not provided and REPO_PATH env var is not set."

    # extract path segment after host
    parsed = urllib.parse.urlparse(api_url)
    url_path = parsed.path.strip("/")
    # take last two segments as search terms (e.g. preovp/preovpdata)
    parts = [p for p in url_path.split("/") if p]
    search_terms = parts[-2:] if len(parts) >= 2 else parts

    results = []

    for term in search_terms:
        # search PHP and Go files
        for pattern in ["**/*.php", "**/*.go"]:
            for f in Path(path).rglob(pattern):
                try:
                    text = f.read_text(errors="ignore")
                except Exception:
                    continue
                if term.lower() in text.lower():
                    # extract up to 80 lines around first match
                    lines = text.splitlines()
                    for i, line in enumerate(lines):
                        if term.lower() in line.lower():
                            start = max(0, i - 5)
                            end = min(len(lines), i + 75)
                            snippet = "\n".join(lines[start:end])
                            rel = str(f).replace(path, "").lstrip("/")
                            results.append(f"### {rel} (line {i+1})\n```\n{snippet}\n```")
                            break  # one snippet per file
                if len(results) >= 6:
                    break
            if len(results) >= 6:
                break

    if not results:
        return f"No source files found matching URL segments: {search_terms} in {path}"

    return "\n\n".join(results[:6])


@mcp.tool()
def query_pg(sql: str, repo_path: str = "") -> str:
    """
    Run a read-only SQL SELECT query on the PostgreSQL database discovered in the repo.
    Use this to fetch valid IDs and values to build API test payloads,
    or to inspect table structure before testing.

    Args:
        sql: A SELECT SQL statement to execute.
        repo_path: Absolute path to the repository root. Defaults to REPO_PATH env var.
    """
    path = repo_path or REPO_PATH
    sql_upper = sql.strip().upper()
    if not sql_upper.startswith("SELECT") and not sql_upper.startswith("WITH"):
        return "ERROR: Only SELECT queries are allowed."
    try:
        conn = _get_pg_conn(path)
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(sql)
            rows = cur.fetchmany(50)  # cap at 50 rows
        conn.close()
        return json.dumps([dict(r) for r in rows], indent=2, default=str)
    except Exception as e:
        return f"ERROR: {e}"


@mcp.tool()
def call_api(
    method: str,
    url: str,
    payload: dict = None,
    extra_headers: dict = None,
) -> str:
    """
    Call an API endpoint using the configured EMP_ID and AK auth credentials.
    Automatically injects empid and AK into the payload for POST requests,
    or as query params for GET requests.

    Args:
        method: HTTP method — GET or POST.
        url: Full URL of the endpoint.
        payload: Dict of body params (for POST). empid and AK are auto-injected.
        extra_headers: Optional additional HTTP headers.
    """
    method = method.upper()
    payload = payload or {}

    # inject auth
    auth = {"empid": EMP_ID, "AK": AK}

    headers = {"Content-Type": "application/x-www-form-urlencoded"}
    if extra_headers:
        headers.update(extra_headers)

    try:
        if method == "POST":
            body = {**auth, **payload}
            data = urllib.parse.urlencode(body).encode()
            req = urllib.request.Request(url, data=data, headers=headers, method="POST")
        else:
            params = {**auth, **payload}
            full_url = url + "?" + urllib.parse.urlencode(params)
            req = urllib.request.Request(full_url, headers=headers, method="GET")

        with urllib.request.urlopen(req, timeout=30) as resp:
            status = resp.status
            body_bytes = resp.read()

        try:
            parsed = json.loads(body_bytes)
            return json.dumps({"status_code": status, "response": parsed}, indent=2)
        except Exception:
            return json.dumps({"status_code": status, "response": body_bytes.decode(errors="replace")}, indent=2)

    except urllib.error.HTTPError as e:
        body = e.read().decode(errors="replace")
        return json.dumps({"status_code": e.code, "error": e.reason, "body": body}, indent=2)
    except Exception as e:
        return f"ERROR: {e}"


@mcp.tool()
def validate_db(sql: str, expected: dict = None, repo_path: str = "") -> str:
    """
    Run a SELECT query to validate database state after an API call.
    Optionally provide an 'expected' dict of field->value pairs to assert against the first row.
    Returns the query result and a PASS/FAIL verdict if expected is provided.

    Args:
        sql: SELECT query to check DB state, e.g. SELECT * FROM orders WHERE id=100
        expected: Optional dict of expected field->value pairs to assert, e.g. {"status": "active"}
        repo_path: Absolute path to the repository root. Defaults to REPO_PATH env var.
    """
    path = repo_path or REPO_PATH
    sql_upper = sql.strip().upper()
    if not sql_upper.startswith("SELECT") and not sql_upper.startswith("WITH"):
        return "ERROR: Only SELECT queries are allowed."
    try:
        conn = _get_pg_conn(path)
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(sql)
            rows = cur.fetchmany(20)
        conn.close()

        result = [dict(r) for r in rows]

        if not expected:
            return json.dumps({"rows": result}, indent=2, default=str)

        if not result:
            return json.dumps({"verdict": "FAIL", "reason": "No rows returned", "rows": []}, indent=2)

        first = result[0]
        mismatches = {}
        for k, v in expected.items():
            actual = first.get(k)
            if str(actual) != str(v):
                mismatches[k] = {"expected": v, "actual": actual}

        if mismatches:
            return json.dumps({"verdict": "FAIL", "mismatches": mismatches, "rows": result}, indent=2, default=str)
        else:
            return json.dumps({"verdict": "PASS", "rows": result}, indent=2, default=str)

    except Exception as e:
        return f"ERROR: {e}"


# ── entrypoint ────────────────────────────────────────────────────────────────

def main():
    mcp.run(transport="stdio")

if __name__ == "__main__":
    main()
