from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from backend.db.session import SessionLocal
from backend.db.models import Universe, Axiom, Theorem
from backend.api.auth import get_api_key
from backend.core.logging_config import get_logger

router = APIRouter()
logger = get_logger("visualization_routes")

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

@router.get("/visualization/universe/{universe_id}")
def get_universe_graph(universe_id: int, db: Session = Depends(get_db), api_key: str = Depends(get_api_key)):
    try:
        universe = db.query(Universe).filter(Universe.id == universe_id).first()
        axioms = db.query(Axiom).filter(Axiom.universe_id == universe_id).all()
        theorems = db.query(Theorem).filter(Theorem.universe_id == universe_id).all()
        nodes = [{"id": ax.id, "type": "axiom", "label": ax.statement} for ax in axioms] + \
                [{"id": th.id, "type": "theorem", "label": th.statement} for th in theorems]
        edges = []
        for th in theorems:
            for ax in axioms:
                if ax.statement in th.proof:
                    edges.append({"source": ax.id, "target": th.id, "type": "proof"})
        logger.info(f"Visualization graph for universe {universe_id} generated.")
        return {
            "universe": {"id": universe.id, "name": universe.name, "type": universe.universe_type},
            "nodes": nodes,
            "edges": edges
        }
    except Exception as e:
        logger.error(f"Visualization error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/visualization/universes")
def get_all_universe_graphs(db: Session = Depends(get_db), api_key: str = Depends(get_api_key)):
    try:
        universes = db.query(Universe).all()
        result = []
        for universe in universes:
            axioms = db.query(Axiom).filter(Axiom.universe_id == universe.id).all()
            theorems = db.query(Theorem).filter(Theorem.universe_id == universe.id).all()
            nodes = [{"id": ax.id, "type": "axiom", "label": ax.statement} for ax in axioms] + \
                    [{"id": th.id, "type": "theorem", "label": th.statement} for th in theorems]
            edges = []
            for th in theorems:
                for ax in axioms:
                    if ax.statement in th.proof:
                        edges.append({"source": ax.id, "target": th.id, "type": "proof"})
            result.append({
                "universe": {"id": universe.id, "name": universe.name, "type": universe.universe_type},
                "nodes": nodes,
                "edges": edges
            })
        logger.info("Visualization graphs for all universes generated.")
        return result
    except Exception as e:
        logger.error(f"Visualization error: {e}")
        raise HTTPException(status_code=500, detail=str(e))