# data/database.py
from sqlalchemy import create_engine, Column, Integer, String, JSON
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from datetime import datetime

Base = declarative_base()

class CompanyData(Base):
    __tablename__ = 'company_data'
    id = Column(Integer, primary_key=True)
    ticker = Column(String)
    data_type = Column(String)
    data = Column(JSON)
    date = Column(String)

class DatabaseManager:
    def __init__(self, db_url='sqlite:///financial_data.db'):
        self.engine = create_engine(db_url)
        self.Session = sessionmaker(bind=self.engine)
        Base.metadata.create_all(self.engine)

    def save_data(self, ticker, data_type, data):
        session = self.Session()
        try:
            new_entry = CompanyData(
                ticker=ticker,
                data_type=data_type,
                data=data,
                date=datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            )
            session.add(new_entry)
            session.commit()
        except Exception as e:
            print(f"Error saving data: {e}")
        finally:
            session.close()

    def get_historical_data(self, ticker):
        session = self.Session()
        try:
            financials = session.query(CompanyData).filter(
                CompanyData.ticker == ticker,
                CompanyData.data_type == 'financials'
            ).order_by(CompanyData.date.desc()).first()

            news = session.query(CompanyData).filter(
                CompanyData.ticker == ticker,
                CompanyData.data_type == 'news'
            ).order_by(CompanyData.date.desc()).all()

            return {
                'financials': financials.data if financials else None,
                'news': [n.data for n in news]
            }
        finally:
            session.close()