from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import Optional, List
import json
import hashlib
import time
import rocksdb
import sqlite3

app = FastAPI()

# === RocksDB setup ===
rocks = rocksdb.DB("rocksdb_data", rocksdb.Options(create_if_missing=True))

# === SQLite setup ===
sqlite_conn = sqlite3.connect("decisions.db")
sqlite_conn.row_factory = sqlite3.Row

# --- Models ---
class BundleRequest(BaseModel):
    start_id: Optional[str] = None
    query: Optional[str] = None
    role: str = Field(..., example="engineer")
    max_hops: int = Field(3, ge=1, le=10)
    top_k: int = Field(10, ge=1, le=50)
    include_full_dr: bool = False
    redaction: bool = True


class DRMetadata(BaseModel):
    id: str
    statement: str
    lifecycle_state: Optional[str]
    role_exposure: dict
    tags: List[str] = []
    timestamp: str
    relation: Optional[str] = None
    score: Optional[float] = None


# --- Helpers ---
def serialize(obj) -> bytes:
    return json.dumps(obj).encode("utf-8")


def deserialize(data: bytes):
    return json.loads(data.decode("utf-8"))


def sha256_of_content(content: str) -> str:
    return hashlib.sha256(content.encode("utf-8")).hexdigest()


def now_iso() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def get_metadata(dr_id: str) -> Optional[dict]:
    # Try RocksDB
    val = rocks.get(f"meta:{dr_id}".encode())
    if val:
        return deserialize(val)
    # Fallback SQLite
    row = sqlite_conn.execute(
        "SELECT statement, lifecycle_state, role_exposure, tags, timestamp FROM decisions WHERE id=?",
        (dr_id,),
    ).fetchone()
    if row:
        return {
            "id": dr_id,
            "statement": row["statement"],
            "lifecycle_state": row["lifecycle_state"],
            "role_exposure": json.loads(row["role_exposure"]),
            "tags": json.loads(row["tags"]),
            "timestamp": row["timestamp"],
        }
    return None


def get_ancestors(dr_id: str) -> List[str]:
    # RocksDB reverse index fallback to SQLite edges
    val = rocks.get(f"rev:{dr_id}".encode())
    if val:
        return deserialize(val)
    rows = sqlite_conn.execute(
        "SELECT source_id FROM edges WHERE target_id=? AND relation='influences'",
        (dr_id,),
    ).fetchall()
    return [r[0] for r in rows]


def walk_back(
    start_id: str, n: int, role: str, top_k: int
) -> List[DRMetadata]:
    from collections import deque
    import heapq

    visited = set()
    results = []
    queue = deque([(start_id, 0)])

    # Dummy embed/sim functions — replace with your actual implementations
    def embed(text): return text
    def cosine_similarity(a, b): return 1.0 if a == b else 0.5

    query_vec = embed("")  # empty here — extend as needed

    while queue:
        dr_id, depth = queue.popleft()
        if dr_id in visited or depth > n:
            continue
        visited.add(dr_id)

        meta = get_metadata(dr_id)
        if not meta:
            continue

        # Role exposure filter
        if role and not meta["role_exposure"].get(role, False):
            continue

        # Score heuristic: favor close ancestors and "approved" states
        dist_score = max(0, (n - depth) / n)
        state_bonus = 1.1 if meta.get("lifecycle_state") == "approved" else 1.0
        score = dist_score * state_bonus

        heapq.heappush(results, (-score, DRMetadata(**meta, score=score)))

        for anc in get_ancestors(dr_id):
            queue.append((anc, depth + 1))

    # Return top_k sorted descending by score
    sorted_results = [md for _, md in heapq.nsmallest(top_k, results)]
    return sorted_results


def assemble_bundle(
    start_id: str,
    role: str,
    max_hops: int,
    top_k: int,
    include_full_dr: bool,
    redact: bool,
    query: Optional[str] = None,
) -> dict:
    # TODO: integrate query embedding and refined scoring later

    timeline = []
    evidence_refs = set()

    ancestors = walk_back(start_id, max_hops, role, top_k)

    for meta in ancestors:
        item = {
            "id": meta.id,
            "timestamp": meta.timestamp,
            "statement": meta.statement if not redact else redact_field(meta.statement, role),
            "lifecycle_state": meta.lifecycle_state,
            "score": meta.score,
            "role_exposure": meta.role_exposure,
            "tags": meta.tags,
            "relation_to_start": meta.relation,
            # "full_payload": ... fetch + decrypt if include_full_dr and authorized
            # For now omitted for brevity
        }
        timeline.append(item)
        # Collect evidence refs from meta if available (stub)
        # evidence_refs.update(meta.get("evidence_refs", []))

    summary = f"Decision bundle starting at {start_id} with {len(timeline)} items."

    bundle_content = {
        "bundle_id": f"bundle:sha256:{sha256_of_content(summary)}",
        "start_id": start_id,
        "generated_at": now_iso(),
        "summary": summary,
        "timeline": timeline,
        "constraints_summary": [],
        "key_evidence_refs": list(evidence_refs),
        "goal_alignment": {},
        "suggested_actions": [],
        "escalation": {"required": False, "who": []},
        "signatures": [],
        "cache_hit": False,
    }

    return bundle_content


def redact_field(text: str, role: str) -> str:
    # Stub: redact sensitive info based on role
    # Replace with your own policy
    if role == "pm":
        return text
    else:
        return text.replace("SECRET", "[REDACTED]")


def cache_get(key: str):
    val = rocks.get(key.encode())
    if val:
        return deserialize(val)
    return None


def cache_set(key: str, value: dict):
    rocks.put(key.encode(), serialize(value))


# --- Endpoint ---
@app.post("/decision/bundle")
async def decision_bundle(req: BundleRequest):
    # Validate inputs
    if not req.start_id and not req.query:
        raise HTTPException(status_code=400, detail="start_id or query required")

    # Resolve anchor if query only
    start_id = req.start_id
    if not start_id:
        # Placeholder: your vector search to find anchor DR by query
        # For now raise error
        raise HTTPException(status_code=400, detail="Query-based anchor resolution not implemented")

    cache_key = f"bundle:{start_id}:{req.role}:{req.max_hops}:{hash(req.query or '')}:{req.top_k}"
    cached_bundle = cache_get(cache_key)
    if cached_bundle:
        cached_bundle["cache_hit"] = True
        return cached_bundle

    bundle = assemble_bundle(
        start_id=start_id,
        role=req.role,
        max_hops=req.max_hops,
        top_k=req.top_k,
        include_full_dr=req.include_full_dr,
        redact=req.redaction,
        query=req.query,
    )

    cache_set(cache_key, bundle)
    return bundle
