import copy
from dataclasses import replace
from datetime import datetime,timedelta,timezone
from pathlib import Path
import tempfile
import unittest
from netplan_v4.domain import Battery,Family,Limits,Price,QuarterPast,Step,default_registry,split_base_load,priced_prefix
from netplan_v4.optimizer import optimize
from netplan_v4.selection import ReplayScore,choose_family,training_due,promote_candidate
from netplan_v4.store import PlannerStore
from netplan_v4.service import ingest,latest,run_once,status,create_app
from fastapi.testclient import TestClient

AT=datetime(2026,10,1,12,0,tzinfo=timezone.utc)
AID='e3a08f9e-af12-4695-99bd-8b51c0520021'
TOKEN='synthetic-test-token-not-real-credentials'

def seq(load,pv=None,buy=None,sell=None,at=AT):
    n=len(load);pv=pv or [0.]*n;buy=buy or [.3]*n;sell=sell or [.1]*n
    return [Step(at+timedelta(minutes=5*i),load[i],pv[i],Price(buy[i]),Price(sell[i])) for i in range(n)]
def batt(**kw):
    values=dict(asset_id='b',capacity_kwh=10.,soc_percent=20.,min_soc_percent=10.,max_soc_percent=90.,max_charge_w=5000.,max_discharge_w=4000.,measured_at=AT,grid_charging=True)
    values.update(kw);return Battery(**values)
def solve(steps,assets=None,**kw):
    kw.setdefault('observed_peaks',{s.start.astimezone(__import__('zoneinfo').ZoneInfo('Europe/Zurich')).strftime('%Y-%m'):5. for s in steps})
    kw.setdefault('peak_prices',{m:0. for m in kw['observed_peaks']})
    return optimize(steps,assets or [],at=steps[0].start,**kw)

def inputs(at=AT):
    env=lambda identifier:{'version':1,'eventId':identifier,'observedAt':at.isoformat()}
    op={**env('op1'),'gridW':-3000.,'meteringBoundary':'common_pcc','batteries':[{'id':'b','capacityKwh':10.,'socPercent':20.,'minSocPercent':10.,'maxSocPercent':90.,'maxChargeW':5000.,'maxDischargeW':4000.,'measuredAt':at.isoformat(),'gridCharging':True}],
        'limits':{'importW':10000.,'exportW':None,'managerMonthLimitsW':{}},'measuredPeaks':{'2026-10':{'kw':5.,'source':'meter_month_register'}},'quarterPast':None}
    points=[{'time':(at+timedelta(minutes=5*i)).isoformat(),'loadW':2000.,'pvW':5000. if i<3 else 0.} for i in range(6)]
    forecast={**env('fc1'),'families':{k:{'loadBasis':'base_load','trainedUntil':(at-timedelta(days=1)).isoformat(),'points':copy.deepcopy(points)} for k in ('3','13','23')}}
    tariff={**env('tar1'),'import':{'mode':'static','tariffId':'buy','staticChfKwh':.3},'export':{'mode':'static','tariffId':'sell','staticChfKwh':.1},'peakChfKwMonth':5.}
    return op,forecast,tariff

class OptimizerTest(unittest.TestCase):
    def test_grid_price_arbitrage(self):
        p=solve(seq([1000.]*6,buy=[.01]*3+[.7]*3),[batt()]);self.assertTrue(p['executable']);self.assertGreater(p['points'][0]['batteryTargetW'],0);self.assertLess(p['points'][-1]['batteryTargetW'],0)
    def test_no_future_energy(self):
        p=solve(seq([3000.]*3+[0.]*3,pv=[0.]*3+[5000.]*3),[batt(soc_percent=10.,grid_charging=False)])
        self.assertTrue(all(x['batteryTargetW']>=-1e-5 for x in p['points'][:3]))
    def test_roundtrip_90_percent(self):
        p=solve(seq([0.]*3+[5000.]*3,pv=[5000.]*3+[0.]*3,buy=[.3]*6,sell=[0.]*6),[batt(soc_percent=10.,grid_charging=False,max_discharge_w=5000.)])
        c=sum(max(0,x['batteryTargetW']) for x in p['points']);d=sum(max(0,-x['batteryTargetW']) for x in p['points']);self.assertAlmostEqual(d/c,.9,places=5)
    def test_nulleinspeisung(self):
        p=solve(seq([0.]*3,pv=[10000.]*3),limits=Limits(export_w=0.));self.assertTrue(p['executable']);self.assertTrue(all(abs(x['gridTargetW'])<1e-4 for x in p['points']))
    def test_nulltariffs(self):
        p=solve(seq([1000.]*3,buy=[0.]*3,sell=[0.]*3));self.assertEqual(p['cashCostChf'],0.)
    def test_negative_prices(self):
        p=solve(seq([1000.]*3,buy=[-.1]*3));self.assertLess(p['energyCostChf'],0.)
    def test_asymmetric_limits_and_balance(self):
        steps=seq([0.]*3+[3000.]*3,pv=[5000.]*3+[0.]*3)
        p=solve(steps,[batt(max_charge_w=700.,max_discharge_w=300.)]);self.assertTrue(p['executable'])
        for s,x in zip(steps,p['points']):
            self.assertLessEqual(x['batteryTargetW'],700.001);self.assertGreaterEqual(x['batteryTargetW'],-300.001)
            self.assertAlmostEqual(x['gridTargetW'],s.residual_w+x['batteryTargetW']+x['pvCurtailmentW'],places=4)
            self.assertTrue(9.999<=x['socEndPercent']['b']<=90.001)
    def test_stale_soc_rejected(self):self.assertFalse(solve(seq([0.]*3),[batt(measured_at=AT-timedelta(hours=1))])['executable'])
    def test_soc_not_fabricated(self):self.assertFalse(solve(seq([0.]*3),[batt(soc_percent=5.)])['executable'])
    def test_peak_increment_only_and_baseline(self):
        p=solve(seq([2000.]*3),observed_peaks={'2026-10':1.5},peak_prices={'2026-10':5.})
        self.assertAlmostEqual(p['additionalPeakCostChf'],2.5);self.assertAlmostEqual(p['baselineCashCostChf'],2.65)
    def test_peak_blocks_unprofitable_charging(self):
        p=solve(seq([1000.]*6,buy=[.01]*3+[.5]*3),[batt(soc_percent=10.)],observed_peaks={'2026-10':1.},peak_prices={'2026-10':20.})
        self.assertLessEqual(p['plannedPeaksKw']['2026-10'],1.0001)
    def test_month_switch(self):
        at=datetime(2026,9,30,21,45,tzinfo=timezone.utc)
        p=solve(seq([2000.]*6,at=at),observed_peaks={'2026-09':2.,'2026-10':.5},peak_prices={'2026-09':5.,'2026-10':5.})
        self.assertAlmostEqual(p['additionalPeakCostChf'],7.5)
    def test_manager_cap_is_not_measured_peak(self):
        p=solve(seq([1000.]*3),limits=Limits(manager_month_limits_w={10:1500.}),observed_peaks={'2026-10':0.},peak_prices={'2026-10':5.})
        self.assertEqual(p['measuredPeaksKw']['2026-10'],0.);self.assertAlmostEqual(p['additionalPeakCostChf'],5.)
    def test_infeasible_not_fake_residual(self):
        p=solve(seq([10000.]*3),limits=Limits(import_w=1000.));self.assertFalse(p['executable']);self.assertEqual(p['points'],[])
    def test_elapsed_quarter_required(self):
        s=seq([1000.],at=AT+timedelta(minutes=10));self.assertFalse(solve(s)['executable'])
        p=solve(s,quarter_history={AT:QuarterPast(.5,600)},peak_prices={'2026-10':5.},observed_peaks={'2026-10':0.})
        self.assertAlmostEqual(p['plannedPeaksKw']['2026-10'],7/3)
    def test_partial_first_step(self):
        s=seq([1000.]*3);s[0]=replace(s[0],start=AT+timedelta(seconds=10),seconds=290)
        p=solve(s,quarter_history={AT:QuarterPast(.02,10)});self.assertTrue(p['executable']);self.assertEqual(p['validFrom'],s[0].start.isoformat())
    def test_terminal_guard(self):
        p=solve(seq([0.]*3,sell=[10.]*3),[batt(soc_percent=50.,grid_charging=False)]);self.assertGreaterEqual(p['points'][-1]['socEndPercent']['b'],49.999)
    def test_grid_charging_opt_in(self):
        p=solve(seq([1000.]*6,buy=[.01]*3+[.5]*3),[batt(soc_percent=10.,grid_charging=False)]);self.assertTrue(all(x['batteryTargetW']<=1e-5 for x in p['points']))

class StoreAndSelectionTest(unittest.TestCase):
    def setUp(self):self.temp=tempfile.TemporaryDirectory();self.store=PlannerStore(str(Path(self.temp.name)/'planner.sqlite'))
    def tearDown(self):self.store.close();self.temp.cleanup()
    def test_settings_revision_and_shadow_only(self):
        s=self.store.save_settings(AID,{'family':'23'},0,AT);self.assertEqual(s['revision'],1)
        with self.assertRaises(ValueError):self.store.save_settings(AID,{'family':'3'},0,AT)
        with self.assertRaises(ValueError):self.store.save_settings(AID,{'runMode':'live'},1,AT)
    def test_request_during_calculation_survives(self):
        self.store.request(AID,'manual',AT);claim=self.store.claim(AT);self.store.request(AID,'prices_changed',AT)
        p=solve(seq([0.]*3),config_revision=0)
        with self.assertRaises(ValueError):self.store.publish_shadow(AID,p,0,AT,claim['sequence'],claim['lease_token'])
        self.store.finish(claim);self.assertIsNotNone(self.store.claim(AT))
    def test_month_peak_persistent_and_monotonic(self):
        self.store.initialize_peak(AID,'2026-10',20.,'meter_month_register',AT)
        with self.assertRaises(ValueError):self.store.initialize_peak(AID,'2026-10',15.,'meter_month_register',AT)
        self.store.close();self.store=PlannerStore(str(Path(self.temp.name)/'planner.sqlite'));self.assertEqual(self.store.peaks(AID)['2026-10'],20.)
    def test_cap_is_not_valid_peak_source(self):
        with self.assertRaises(ValueError):self.store.initialize_peak(AID,'2026-10',20.,'configured_limit',AT)
    def test_meter_peak_requires_complete_quarter(self):
        self.store.initialize_peak(AID,'2026-10',0.,'verified_new_month',AT)
        for i in range(3):r=self.store.record_import_interval(AID,AT+timedelta(minutes=5*i),.5,AT+timedelta(hours=1))
        self.assertEqual(r['quarterPeakKw'],6.);self.assertEqual(self.store.peaks(AID)['2026-10'],6.)
    def test_shadow_ack_cannot_claim_applied(self):
        p=solve(seq([0.]*3),config_revision=0);self.store.publish_shadow(AID,p,0,AT)
        with self.assertRaises(ValueError):self.store.acknowledge(AID,p['planId'],0,AT,status='applied')
        self.store.acknowledge(AID,p['planId'],0,AT,status='shadow_seen')
    def test_family_registry_extensible(self):
        r=default_registry();r.register(Family('42','p42','l42','g42','Future'));self.assertEqual(r.get('42').pv,'p42')
    def test_sdl_removed_once(self):
        self.assertEqual(split_base_load(15000,[3000],2000),10000);self.assertEqual(split_base_load(13000,[3000],2000,True),10000)
        with self.assertRaises(ValueError):split_base_load(100,[1000],0)
    def test_auto_uses_comparable_cost_and_margin(self):
        scores=[ReplayScore(k,'same',AT-timedelta(days=14),AT,AT-timedelta(days=14),AT,c,1.,14) for k,c in [('3',10.),('13',9.5),('23',5.)]]
        self.assertEqual(choose_family('auto','3',scores,now=AT)['family'],'23');self.assertEqual(choose_family('auto','3',scores,now=AT,margin_chf=6.)['family'],'3')
    def test_auto_no_future_outcomes(self):
        scores=[ReplayScore(k,'same',AT-timedelta(days=14),AT,AT-timedelta(days=14),AT+timedelta(hours=1),c,1.,14) for k,c in [('3',10.),('13',9.5),('23',5.)]]
        self.assertEqual(choose_family('auto','3',scores,now=AT)['mode'],'collecting')
    def test_training_schedule_and_validation(self):
        self.assertTrue(training_due(AT-timedelta(days=1),AT));self.assertFalse(training_due(AT-timedelta(days=1),AT,'weekly'))
        self.assertFalse(promote_candidate(active_cost=10,candidate_cost=1,valid_coverage=True,no_data_leakage=False,constraints_passed=True))

class ServiceIntegrationTest(unittest.TestCase):
    def setUp(self):self.temp=tempfile.TemporaryDirectory();self.store=PlannerStore(str(Path(self.temp.name)/'planner.sqlite'))
    def tearDown(self):self.store.close();self.temp.cleanup()
    def load(self,values=None,now=AT):
        values=values or inputs()
        for kind,value in zip(('operation','forecast','tariffs'),values):ingest(self.store,AID,kind,value,now)
    def test_pipeline_shadow_plan(self):
        self.load();p=run_once(self.store,AT);self.assertTrue(p['executable']);self.assertEqual(p['runMode'],'shadow');self.assertTrue(status(self.store,AID,AT)['fresh']);self.assertEqual(self.store.current(AID)['planId'],p['planId'])
    def test_exact_model_selection(self):
        self.load();self.store.save_settings(AID,{'family':'23'},0,AT);self.assertEqual(run_once(self.store,AT)['sourceFamily'],'23')
    def test_one_tick_not_every_poll(self):
        self.load();run_once(self.store,AT);self.assertEqual(run_once(self.store,AT)['status'],'idle')
    def test_stale_operation_visible(self):
        self.load();p=run_once(self.store,AT+timedelta(minutes=5));self.assertEqual(p['status'],'awaiting_inputs');self.assertIsNone(self.store.current(AID))
    def test_missing_peak_not_invented(self):
        op,fc,tar=inputs();op['measuredPeaks']={};self.load((op,fc,tar));self.assertEqual(run_once(self.store,AT)['status'],'awaiting_inputs')
    def test_static_zero_roundtrip(self):
        op,fc,tar=inputs();tar['import']['staticChfKwh']=0.;tar['export']['staticChfKwh']=0.;tar['peakChfKwMonth']=0.;self.load((op,fc,tar));self.assertEqual(run_once(self.store,AT)['cashCostChf'],0.)
    def test_immutable_id_conflict(self):
        op,fc,tar=inputs();self.load((op,fc,tar));self.assertEqual(ingest(self.store,AID,'operation',op,AT)['status'],'duplicate');op['gridW']=123.
        with self.assertRaises(ValueError):ingest(self.store,AID,'operation',op,AT)
    def test_older_input_never_overwrites_current(self):
        op,fc,tar=inputs();self.load((op,fc,tar));op['observedAt']=(AT-timedelta(seconds=1)).isoformat();op['eventId']='old'
        self.assertEqual(ingest(self.store,AID,'operation',op,AT)['status'],'archived_older');self.assertEqual(latest(self.store,AID,'operation')['eventId'],'op1')
    def test_unpublished_dynamic_tail_not_executed(self):
        op,fc,tar=inputs();tar['import']={'mode':'dynamic','tariffId':'buy'};self.load((op,fc,tar));self.assertEqual(run_once(self.store,AT)['status'],'awaiting_inputs')
    def test_published_negative_price_and_horizon(self):
        op,fc,tar=inputs();tar['import']={'mode':'dynamic','tariffId':'buy'};self.load((op,fc,tar))
        price={'version':1,'eventId':'price1','observedAt':AT.isoformat(),'periods':[{'tariffId':'buy','side':'import','start':AT.isoformat(),'end':(AT+timedelta(minutes=15)).isoformat(),'value':-2.,'unit':'Rp/kWh','observedAt':AT.isoformat(),'sourceKind':'published_interval'}]}
        ingest(self.store,AID,'prices',price,AT);p=run_once(self.store,AT);self.assertEqual(len(p['points']),3);self.assertEqual(p['points'][0]['importPriceChfKwh'],-.02)
    def test_carried_price_does_not_extend_horizon(self):
        op,fc,tar=inputs();tar['import']={'mode':'dynamic','tariffId':'buy'};self.load((op,fc,tar))
        price={'version':1,'eventId':'price1','observedAt':AT.isoformat(),'periods':[{'tariffId':'buy','side':'import','start':AT.isoformat(),'end':(AT+timedelta(hours=48)).isoformat(),'value':.1,'unit':'CHF/kWh','observedAt':AT.isoformat(),'sourceKind':'carried_forward'}]}
        ingest(self.store,AID,'prices',price,AT);self.assertEqual(run_once(self.store,AT)['status'],'awaiting_inputs')
    def test_aggregate_load_warning(self):
        op,fc,tar=inputs();fc['families']['3']['loadBasis']='house_total';self.load((op,fc,tar));self.assertTrue(run_once(self.store,AT)['warnings'])
    def test_extra_credentials_not_stored(self):
        op,fc,tar=inputs();op['password']='do-not-store'
        with self.assertRaises(ValueError):ingest(self.store,AID,'operation',op,AT)
    def test_http_auth_allowlist_and_revision_conflict(self):
        app=create_app(str(Path(self.temp.name)/'http.sqlite'),TOKEN,[AID],start_worker=False)
        with TestClient(app) as client:
            path=f'/internal/v2/prognosis/{AID}/planner';headers={'X-Enelix-Service-Token':TOKEN}
            self.assertEqual(client.get(path).status_code,401);self.assertEqual(client.get(path,headers=headers).status_code,200)
            other='30509683-4569-49e4-848f-4905e4cc813a';self.assertEqual(client.get(path.replace(AID,other),headers=headers).status_code,403)
            self.assertEqual(client.put(path+'/settings',headers=headers,json={'expectedRevision':0,'changes':{'family':'23'}}).status_code,200)
            self.assertEqual(client.put(path+'/settings',headers=headers,json={'expectedRevision':0,'changes':{'family':'3'}}).status_code,409)
            self.assertEqual(client.put(path+'/settings',headers=headers,json={'expectedRevision':1,'changes':{'runMode':'live'}}).status_code,400)
    def test_body_limit_and_no_legacy_db(self):
        with self.assertRaises(ValueError):create_app(str(Path(self.temp.name)/'users.db'),TOKEN,[AID])
        app=create_app(str(Path(self.temp.name)/'http.sqlite'),TOKEN,[AID],start_worker=False)
        with TestClient(app) as client:
            r=client.post(f'/internal/v2/prognosis/{AID}/planner/inputs/forecast',json={'x':'a'*2000001});self.assertEqual(r.status_code,413)

if __name__=='__main__':unittest.main()